Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

Python으로 해결하는 공정한 사탕 교환(Fair Candy Swap) 알고리즘

문제 소개

A와 B라는 두 친구가 서로 다른 크기의 사탕 바를 가지고 있다고 가정해 봅시다. 여기서 A[i]는 A가 가진 i번째 사탕 바의 크기를, B[j]는 B가 가진 j번째 사탕 바의 크기를 의미합니다.

두 사람은 친구이기 때문에 각자 사탕 바 하나씩을 서로 교환한 뒤, 두 사람이 가진 사탕의 총량(가진 사탕 바 크기들의 합)이 완전히 같아지도록 만들고 싶어 합니다. 따라서 우리는 정수 배열 ans를 반환해야 하며, ans[0]에는 A가 내놓아야 할 사탕 바의 크기를, ans[1]에는 B가 내놓아야 할 사탕 바의 크기를 담으면 됩니다. 조건을 만족하는 답이 여러 개라면 그중 하나만 반환하면 됩니다.

예를 들어 A = [1, 2], B = [2, 3]일 때 출력 결과는 [1, 2]가 됩니다.

문제 해결 접근 방법

교환 전후의 사탕 총량을 수식으로 표현하면 다음 관계가 성립합니다.

(sum(A) - x + y) == (sum(B) - y + x)

이 식을 정리하면 x - y = (sum(A) - sum(B)) / 2가 되므로, A에서 내놓을 사탕과 B에서 내놓을 사탕의 크기 차이는 항상 일정한 값(diff)이 됩니다. 이 원리를 활용하면 문제를 효율적으로 해결할 수 있습니다.

  • A의 합과 B의 합의 차이를 구한 뒤 2로 나누고, 그 정수 몫을 diff 변수에 저장합니다.
  • B를 집합(set) 자료형으로 변환하면 특정 값이 존재하는지 O(1) 시간에 확인할 수 있어 탐색 속도가 크게 향상됩니다.
  • A의 각 원소 i에 대해 다음을 반복합니다.
    • i − diff가 집합 B 안에 존재한다면, 즉시 [i, i − diff]를 반환합니다.

Python 코드 구현

아래 예시 코드를 통해 더 자세히 이해해 보겠습니다.

class Solution(object):
   def fairCandySwap(self, A, B):
      diff = (sum(A) - sum(B))//2
      B=set(B)
      for i in A:
         if i- diff in B:
            return [i,i-diff]
ob1 = Solution()
print(ob1.fairCandySwap([1,2], [2,3]))

입력

[1,2]
[2,3]

출력

[1,2]

복잡도 분석

이 알고리즘은 A와 B의 모든 원소를 한 번씩 순회하므로 시간 복잡도는 O(n + m), 집합 변환에 따른 추가 공간 때문에 공간 복잡도 역시 O(m)입니다(n은 A의 길이, m은 B의 길이). 덕분에 무차별 대입 방식(O(n × m))보다 훨씬 효율적으로 동작합니다.