두 배열의 원소 교체

이것이 취업을 위한 코딩 테스트다 with 파이썬을 공부하면서 정리한 내용입니다.

---

문제 내용

• 두 배열 A와 B는 N개의 원소로 구성되어 있으며, 배열의 원소는 모두 자연수
• 최대 K번의 바꿔치기 연산을 수행할 수 있는데, 바꿔치기 연산이란 배열 A에 있는 원소 하나와 배열 B에 있는 원소 하나를 골라서 두 원소를 서로 바꾸는 것
• 최종 목표는 배열 A의 모든 원소의 합이 최대가 되도록 하는 것
• N, K, 그리고 배열 A와 B의 정보가 주어졌을 때, 최대 K번의 바꿔치기 연산을 수행하여 만들 수 있는 배열 A의 모든 원소의 합의 최댓값을 출력하는 프로그램을 작성

입력 조건

• 첫 번째 줄에 N, K가 공백으로 입력됨 $(1 \le N \le 100,000)$ $(0 \le K \le N)$
• 두 번째 줄에 배열 A의 원소들이 공백으로 구분되어 입력됨
• 세 번째 줄에 배열 B의 원소들이 공백으로 구분되어 입력됨
• 배열의 모든 원소는 10,000,000보다 작은 자연수

출력 조건

• 최대 K번의 바꿔치기 연산을 수행하여 만들 수 있는 배열 A의 모든 원소의 합의 최댓값을 출력

입력 예시

5 3
1 2 5 4 3
5 5 6 6 5

출력 예시

26

---

문제 해설

• 문제를 해결하기 위한 기본 아이디어는 매번 배열 A에서 가장 작은 원소를 골라서, 배열 B에서 가장 큰 원소와 교체를 하는 것
• 단, 배열 A에서 가장 작은 원소가 배열 B에서 가장 큰 원소보다 작을 때에만 교체를 수행
• 이러한 과정을 K번 반복
• 배열 A를 오름차순으로 정렬하고 배열 B를 내림차순으로 정렬
• 두 배열의 원소를 가장 첫 번째 인덱스부터 차례대로 비교
• 두 배열의 원소가 최대 100,000개까지 입력될 수 있으므로 $O(NlogN)$을 보장하는 정렬 알고리즘을 이용

소스 코드

n, k = map(int, input().split())     # N과 K를 입력받기
a = list(map(int, input().split()))  # 배열 A의 모든 원소 입력받기
b = list(map(int, input().split()))  # 배열 B의 모든 원소 입력받기

a.sort() # 배열 A는 오름차순 정렬 수행
b.sort(reverse=True) # 배열 B는 내림차순 정렬 수행

# 첫 번째 인덱스부터 확인하며, 두 배열의 원소를 최대 K번 비교
for i in range(k):
    # A의 원소가 B의 원소보다 작은 경우
    if a[i] < b[i]:
        # 두 원소를 교체
        a[i], b[i] = b[i], a[i]
    else: # A의 원소가 B의 원소보다 크거나 같을 때, 반복문을 탈출
        break

print(sum(a)) # 배열 A의 모든 원소의 합을 출력

---

나의 풀이

• 문제 해설과 동일

소스 코드

import sys

n, k = map(int, sys.stdin.readline().rstrip().split())
a = list(map(int, sys.stdin.readline().rstrip().split()))
b = list(map(int, sys.stdin.readline().rstrip().split()))

a.sort()
b.sort(reverse=True)

for i in range(k):
    if a[i] < a[i]:
        a[i], b[i] = b[i], a[i]
    else:
        break

print(sum(a))