Python으로 두 배열의 합을 같게 만드는 최소 연산 횟수 구하기

문제 설명

두 개의 리스트 nums1과 nums2가 주어지며, 두 리스트의 모든 요소는 1부터 6 사이의 값을 가집니다. 사용할 수 있는 연산은 nums1 또는 nums2에서 숫자 하나를 골라 그 값을 1부터 6 사이의 임의의 숫자로 바꾸는 것입니다. 이때 두 배열의 합이 서로 같아지도록 만드는 데 필요한 최소 연산 횟수를 구해야 하며, 어떻게 해도 같게 만들 수 없다면 -1을 반환합니다.

예를 들어 nums1 = [1, 4], nums2 = [5, 4, 4]가 입력으로 주어졌다고 해봅시다. 먼저 nums1의 1을 6으로 바꾸면 nums1의 합은 10이 되고, 이어서 nums2의 4 하나를 1로 바꾸면 nums2의 합도 10이 됩니다. 따라서 정답은 2입니다.

접근 방법: 탐욕 알고리즘

이 문제는 탐욕(Greedy) 기법으로 효율적으로 해결할 수 있습니다. 핵심 아이디어는 한 번의 연산으로 얻을 수 있는 합의 변화량이 가장 큰 원소부터 차례대로 처리하는 것입니다.

1. 합 계산: nums1의 합을 sa, nums2의 합을 sb로 구합니다.
2. 배열 교환: sa > sb라면 nums1과 nums2, 그리고 sa와 sb를 서로 교환합니다. 이렇게 하면 항상 nums1의 합이 더 작거나 같아지므로, 남은 작업은 nums1의 합을 올리거나 nums2의 합을 내리는 것뿐입니다.
3. 정렬: nums1은 오름차순으로 정렬합니다. 값이 작을수록 6으로 바꿨을 때 증가폭(6 − 값)이 크기 때문입니다. 반대로 nums2는 내림차순으로 정렬합니다. 값이 클수록 1로 바꿨을 때 감소폭(값 − 1)이 크기 때문입니다.
4. 초기화: res := 0, toadd := sb − sa, 인덱스 i := 0, j := 0으로 설정합니다.
5. 반복 처리: toadd > 0인 동안 다음을 수행합니다.
   • res를 1 증가시킵니다.
   • i와 j가 각각 nums1, nums2의 범위 안에 있다면 resa = 6 − nums1[i], resb = nums2[j] − 1을 계산하고, 효과가 더 큰 쪽을 적용합니다. resa > resb면 toadd에서 resa를 빼고 i를 증가시키고, 그렇지 않으면 toadd에서 resb를 빼고 j를 증가시킵니다.
   • i만 유효하다면 resa = 6 − nums1[i]를 toadd에서 빼고 i를 증가시킵니다.
   • j만 유효하다면 resb = nums2[j] − 1을 toadd에서 빼고 j를 증가시킵니다.
   • 두 인덱스 모두 범위를 벗어났다면 더 이상 조정할 수 없으므로 -1을 반환합니다.
6. 결과 반환: 반복이 종료되면 res를 반환합니다.

구현 예시

아래 파이썬 코드를 통해 전체 로직을 확인할 수 있습니다.

def solve(nums1, nums2):
    sa = sum(nums1)
    sb = sum(nums2)
    if sa > sb:
        nums1, nums2 = nums2, nums1
        sa, sb = sb, sa

    nums1.sort()
    nums2.sort(reverse=True)
    res = 0
    toadd = sb - sa
    i = 0
    j = 0
    while toadd > 0:
        res += 1
        if i < len(nums1) and j < len(nums2):
            resa = 6 - nums1[i]
            resb = nums2[j] - 1
            if resa > resb:
                toadd -= resa
                i += 1
            else:
                toadd -= resb
                j += 1
        elif i < len(nums1):
            resa = 6 - nums1[i]
            toadd -= resa
            i += 1
        elif j < len(nums2):
            resb = nums2[j] - 1
            toadd -= resb
            j += 1
        else:
            return -1

    return res

nums1 = [1, 4]
nums2 = [5, 4, 4]
print(solve(nums1, nums2))

실행 결과

nums1의 합은 5, nums2의 합은 13이므로 두 합의 차이는 8입니다. 첫 번째 연산에서 nums1의 1을 6으로 바꿔 차이를 5만큼 줄이고(toadd = 3), 두 번째 연산에서 nums2의 5를 1로 바꿔 나머지 4를 처리하면 두 배열의 합이 같아집니다. 따라서 결과는 다음과 같습니다.

2

복잡도 분석

정렬에 O(n log n)의 시간이 걸리고, 이후의 탐욕 순회는 각 원소를 한 번씩만 방문하므로 O(n)입니다. 따라서 전체 시간 복잡도는 O(n log n)이며, 추가 공간 복잡도는 O(1)입니다.