문제 소개
두 개의 배열 nums1과 nums2가 주어졌을 때, 아래 두 규칙 중 하나를 만족하는 삼중항(triplet)의 총 개수를 구하는 프로그램을 만들어 보겠습니다.
- 타입 1: nums1[i]^2 = nums2[j] * nums2[k] 를 만족하는 삼중항 (i, j, k) — 단, 0 <= i < len(nums1) 이고 0 <= j < k < len(nums2)
- 타입 2: nums2[i]^2 = nums1[j] * nums1[k] 를 만족하는 삼중항 (i, j, k) — 단, 0 <= i < len(nums2) 이고 0 <= j < k < len(nums1)
예를 들어 nums1 = [7, 4], nums2 = [5, 2, 8, 9]라고 하면 정답은 1입니다. 타입 1의 삼중항 (1, 1, 2)에서 nums1[1]^2 = 16 = nums2[1] * nums2[2] = 2 * 8 이 성립하기 때문입니다.
풀이 접근 방식
핵심 아이디어는 각 숫자의 등장 횟수를 먼저 집계한 뒤, "한 배열의 어떤 수 t의 제곱(t^2)이 다른 배열의 두 수의 곱으로 표현되는 경우"를 효율적으로 세는 것입니다.
- Counter를 사용해 nums1, nums2의 각 원소와 그 개수를 각각 cnt1, cnt2에 저장합니다.
- triplets(arr1, arr2) 함수를 정의합니다. 이 함수는 "arr1의 어떤 수의 제곱 = arr2의 두 수의 곱"인 경우의 수를 반환합니다.
- arr1의 각 원소 t(등장 횟수 v)에 대해 다음을 수행합니다.
- k := arr2에 t가 존재하면 그 개수, 없으면 0
- tmp := k * (k - 1) / 2 — arr2에서 같은 값 t 두 개를 뽑아 곱했을 때 t^2이 되는 조합의 수
- sq := t * t
- arr2의 각 원소 m에 대해, m < t 이고 sq가 m으로 나누어떨어지면 tmp에 arr2[m] * arr2[sq / m]을 더합니다. 즉, m × (sq / m) = t^2 을 만족하는 쌍의 개수를 누적합니다.
- ans에 tmp * v를 더하고, 모든 반복이 끝나면 ans를 반환합니다.
- 메인에서는 방향이 반대인 경우도 세야 하므로 triplets(cnt1, cnt2) + triplets(cnt2, cnt1)을 반환합니다.
m < t 조건을 사용하면 각 약수 쌍을 한 번만 세게 되므로 중복 계산을 피할 수 있습니다.
예제 코드
다음 구현을 통해 더 잘 이해해 보겠습니다.
from collections import Counter
def solve(nums1, nums2):
cnt1 = Counter(nums1)
cnt2 = Counter(nums2)
def triplets(arr1, arr2):
ans = 0
for t, v in arr1.items():
# arr2에서 같은 값 t 두 개로 t*t를 만드는 경우의 수
k = arr2.get(t, 0)
tmp = k * (k - 1) // 2
sq = t * t
# m * (sq // m) == sq 를 만족하는 쌍 찾기
for m in arr2:
if m < t and sq % m == 0:
tmp += arr2.get(m, 0) * arr2.get(sq // m, 0)
ans += tmp * v
return ans
return triplets(cnt1, cnt2) + triplets(cnt2, cnt1)
nums1 = [7, 4]
nums2 = [5, 2, 8, 9]
print(solve(nums1, nums2))
입력
[7,4], [5,2,8,9]
출력
1
이 풀이에서 시간 복잡도는 서로 다른 원소의 개수를 각각 N, M이라 할 때 대략 O(N * M)이며, 공간 복잡도는 카운터 저장을 위해 O(N + M)입니다. 배열의 크기가 수천 수준까지라면 충분히 빠르게 동작합니다.