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

파이썬으로 배열에서 '좋은 쌍(Nice Pair)' 개수 계산하기

문제 개요

음이 아닌 정수로 구성된 배열 nums가 주어졌을 때, 배열에 존재하는 '좋은 쌍(nice pair)'의 개수를 구하는 프로그램을 작성해 보겠습니다. 답이 매우 커질 수 있으므로, 최종 결과는 10^9 + 7로 나눈 나머지를 반환해야 합니다.

인덱스 쌍 (i, j)가 다음 두 조건을 모두 만족할 때 '좋은 쌍'이라고 정의합니다.

  • 0 <= i < j < nums의 길이
  • nums[i] + rev(nums[j]) == nums[j] + rev(nums[i])

참고: rev() 함수는 정수의 숫자만 뒤집습니다. 예를 들어 rev(564)는 465이며, rev(540)는 뒤집으면 045가 되지만 선행 0이 제거되어 45를 반환합니다.

예시

입력이 nums = [97, 2, 42, 11]이라면 출력은 2입니다. (0, 2)와 (1, 3) 두 쌍이 좋은 쌍에 해당하기 때문입니다.

  • (0, 2): 97 + rev(42) = 97 + 24 = 121, 그리고 42 + rev(97) = 42 + 79 = 121 → 성립
  • (1, 3): 2 + rev(11) = 2 + 11 = 13, 그리고 11 + rev(2) = 11 + 2 = 13 → 성립

접근 방법

이 문제의 핵심은 조건식을 변형하는 것입니다. nums[i] + rev(nums[j]) == nums[j] + rev(nums[i]) 식을 정리하면 nums[i] − rev(nums[i]) == nums[j] − rev(nums[j])가 됩니다. 즉, 각 숫자에서 '자기 자신을 뒤집은 값'을 뺀 차이가 동일한 숫자들끼리만 좋은 쌍을 형성할 수 있습니다.

따라서 해시맵을 이용해 동일한 차이값을 가지는 원소들을 그룹화한 뒤, 각 그룹에서 두 개를 선택하는 조합의 수를 모두 더하면 됩니다. 같은 차이값을 가진 원소가 val개라면, 해당 그룹에서 만들 수 있는 쌍의 개수는 val × (val − 1) / 2입니다.

알고리즘 단계

  • m := 10^9 + 7
  • dic := 기본값이 0인 빈 맵(defaultdict)
  • nums의 각 num에 대해:
    • rev := num의 자릿수를 뒤집은 값
    • dic[num − rev] 값을 1 증가
  • res := 0
  • dic의 모든 값 val에 대해: res := res + (val × (val − 1)) // 2
  • res mod m 반환

이 방식은 각 원소를 한 번씩만 순회하므로 전체 시간 복잡도는 O(n)입니다. 모든 쌍을 일일이 확인하는 브루트 포스 방식(O(n²))보다 훨씬 효율적입니다.

구현 예제

다음 구현을 통해 더 잘 이해해 보겠습니다.

from collections import defaultdict
def solve(nums):
   m = (10**9)+7
   dic = defaultdict(int)
   for num in nums:
      rev=int(str(num)[::-1])
      dic[num-rev]+=1

   res=0
   for val in dic.values():
      res += (val*(val-1)) // 2

   return res % m

nums = [97,2,42,11]
print(solve(nums))

입력

[97,2,42,11]

출력

2