파이썬으로 앞쪽 두 수의 곱과 뒤쪽 두 수의 곱이 같은 쿼드러플 개수 구하기

문제 설명

서로 다른 양의 정수로만 이루어진 리스트 nums가 있다고 가정해 보겠습니다. 이때 nums에서 네 개의 서로 다른 원소 (a, b, c, d)를 골라 a × b = c × d를 만족하는 쿼드러플(quadruple)이 총 몇 개 있는지 찾아야 합니다.

예를 들어 입력이 nums = [3, 6, 4, 8]이라면 출력은 8이 됩니다. 조건을 만족하는 쿼드러플은 [[3,8,6,4], [3,8,4,6], [8,3,6,4], [8,3,4,6], [6,4,3,8], [4,6,3,8], [6,4,8,3], [4,6,8,3]]로 총 8개이기 때문입니다.

해결 접근 방법

모든 가능한 네 원소 조합을 일일이 확인하면 시간이 오래 걸립니다. 대신 두 원소의 곱에 주목하면 효율적으로 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

먼저 리스트에서 만들 수 있는 모든 두 원소 쌍의 곱을 계산하여 해시 맵(딕셔너리)에 빈도수와 함께 저장합니다.
같은 곱을 가진 쌍이 k개 있다면, 서로 다른 두 쌍을 순서대로 고르는 경우의 수는 k × (k − 1)입니다.
각 쌍 내부에서 두 원소의 순서를 바꿀 수 있으므로 (a, b)와 (b, a), (c, d)와 (d, c) 각각 2가지씩, 즉 총 4배를 곱해 최종 답을 구합니다.

이를 단계별로 정리하면 다음과 같습니다.

c := 새로운 딕셔너리(맵)
n := nums의 크기
i를 0부터 n−1까지 반복:
- j를 i+1부터 n−1까지 반복:
  - x := nums[i] × nums[j]
  - c[x] := 기존 값(없으면 0) + 1
ret := 0
c의 모든 값 x에 대해:
- ret := ret + x × (x − 1)
ret × 4를 반환

예제 코드

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

def solve(nums):
    c = {}
    n = len(nums)
    for i in range(n):
        for j in range(i + 1, n):
            x = nums[i] * nums[j]
            c[x] = c.get(x, 0) + 1
    ret = 0
    for x in c.values():
        ret += x * (x - 1)
    return ret * 4

nums = [3, 6, 4, 8]
print(solve(nums))

입력

[3, 6, 4, 8]

출력

동작 원리 살펴보기

입력 [3, 6, 4, 8]에서 만들 수 있는 쌍의 곱은 다음과 같습니다.

3 × 6 = 18
3 × 4 = 12
3 × 8 = 24
6 × 4 = 24
6 × 8 = 48
4 × 8 = 32

여기서 곱이 24인 쌍이 2개이므로, 2 × (2 − 1) = 2가 되고, 여기에 4를 곱한 8이 최종 결과입니다. 나머지 곱들은 각각 한 번씩만 등장하므로 기여하지 않습니다.

시간 복잡도

모든 쌍을 한 번씩 확인하므로 시간 복잡도는 O(n²)이며, 추가로 사용되는 공간은 곱의 종류 수에 비례하므로 최악의 경우 O(n²)입니다. 완전 탐색으로 네 원소를 직접 고르는 O(n⁴) 방식보다 훨씬 효율적입니다.