[파이썬 알고리즘] 두 배열의 원소 교체 문제 풀이

❓ 문제

---

✏ 나의 풀이

---

_, k = map(int, input().split())
a = sorted(map(int, input().split()))
b = sorted(map(int, input().split()))

a[:k] = b[-k:]

print(sum(a))

💡 문제 해결 아이디어

---

매번 배열 A에서 가장 작은 원소를 골라 배열 B에서 가장 큰 원소와 교체하면 된다.

배열 A는 오름차순, 배열 B는 내림차순으로 정렬한다.

이후 두 배열의 원소를 첫 번째 인덱스부터 차례대로 확인하며 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의 모든 원소의 합을 출력

💬 풀이 평가

---

알고리즘을 조금 더 활용하면 좋을 것 같다.

Notice: 이 게시물은 한빛미디어의 이것이 코딩 테스트다 교재를 참고하였습니다.