파이썬으로 특정 범위 내 XOR 쌍의 개수를 구하는 프로그램

문제 설명

배열 nums와 두 값 l, r이 주어졌을 때, '좋은 쌍(nice pair)'의 개수를 구하는 프로그램을 만들어 보겠습니다. 여기서 좋은 쌍이란 인덱스 쌍 (i, j) 중에서 0 <= i < j < len(nums)이고, l <= (nums[i] XOR nums[j]) <= r 조건을 만족하는 쌍을 의미합니다.

예를 들어 입력이 nums = [4, 1, 7, 2], l = 2, r = 6이라면 출력은 6이 됩니다. 조건을 만족하는 좋은 쌍은 다음과 같습니다.

(0, 1): 4 XOR 1 = 5
(0, 2): 4 XOR 7 = 3
(0, 3): 4 XOR 2 = 6
(1, 2): 1 XOR 7 = 6
(1, 3): 1 XOR 2 = 3
(2, 3): 7 XOR 2 = 5

해결 접근 방법

모든 쌍을 하나씩 확인하는 O(n²) 완전 탐색 대신, 비트(bit) 단위로 처리하면 더 효율적으로 문제를 해결할 수 있습니다. 핵심 아이디어는 'XOR 결과가 x 미만인 쌍의 개수'를 반환하는 함수 test(nums, x)를 정의한 뒤, 최종 답을 test(nums, r + 1) − test(nums, l)로 계산하는 것입니다. 이렇게 하면 범위 [l, r]에 포함되는 쌍의 개수를 깔끔하게 얻을 수 있습니다.

test() 함수의 동작 단계는 다음과 같습니다.

count := nums에 있는 각 요소의 빈도수를 저장한 맵(Counter)
res := 0
x가 0이 아닌 동안 반복합니다.
- x가 홀수라면, res += Σ(count[a] × count[(x − 1) XOR a])를 누적합니다.
- count := 키를 a >> 1로, 값을 count[a] + count[a ^ 1]로 갱신한 새로운 맵
- x := x >> 1 (x를 2로 나눈 몫)
res // 2를 반환합니다. (각 쌍이 두 번씩 세어지므로 2로 나눕니다.)

예제 코드

다음 구현을 통해 동작 방식을 더 자세히 이해해 보겠습니다.

from collections import Counter

def solve(nums, l, r):
    def test(nums, x):
        count = Counter(nums)
        res = 0
        while x:
            if x & 1:
                res += sum(count[a] * count[(x - 1) ^ a] for a in count)
            count = Counter({a >> 1: count[a] + count[a ^ 1] for a in count})
            x >>= 1
        return res // 2

    return test(nums, r + 1) - test(nums, l)

nums = [4, 1, 7, 2]
l = 2
r = 6
print(solve(nums, l, r))

입력

[4,1,7,2], 2, 6