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

Python으로 배열 요소의 최대 XOR 찾기: 쿼리 제한 조건 기반 효율적 알고리즘

문제 설명

음이 아닌 정수로 구성된 배열 nums와 쿼리 배열 queries가 주어집니다. 각 쿼리 queries[i]는 (xi, mi) 형태의 쌍으로 이루어져 있으며, i번째 쿼리의 답은 nums의 원소 중 mi 이하인 값과 xi를 XOR 연산했을 때 얻을 수 있는 최댓값입니다. 만약 nums의 모든 원소가 mi보다 크다면 해당 쿼리의 답은 -1이 됩니다. 즉, 쿼리 개수와 같은 크기의 답 배열을 만들어 각 쿼리의 결과를 순서대로 반환해야 합니다.

예를 들어 입력이 다음과 같다고 가정해 보겠습니다.

nums = [0,1,2,3,4], queries = [[3,1],[1,3],[5,6]]

이 경우 출력은 [3, 3, 7]이 됩니다. 그 이유는 다음과 같습니다.

  • 첫 번째 쿼리 [3, 1]: 1 이하인 원소는 0과 1입니다. 0 XOR 3 = 3, 1 XOR 3 = 2이므로 더 큰 값인 3이 답입니다.
  • 두 번째 쿼리 [1, 3]: 3 이하인 원소 중 1 XOR 2 = 3이 최댓값입니다.
  • 세 번째 쿼리 [5, 6]: 6 이하인 원소 중 5 XOR 2 = 7이 최댓값입니다.

해결 접근 방법

이 문제는 비트 단위 접두사(prefix) 집합을 활용하는 그리디 방식으로 효율적으로 해결할 수 있습니다. 핵심 아이디어는 쿼리를 제한값(limit) 기준으로 정렬한 뒤, 상위 비트부터 차례로 검사하면서 각 비트를 가능한 한 1로 만드는 것입니다. 구체적인 알고리즘은 다음과 같습니다.

  1. m := nums의 크기, n := queries의 크기로 설정합니다.
  2. 각 쿼리를 (인덱스 i, x, limit) 형태의 튜플로 변환한 후, limit 값을 기준으로 오름차순 정렬합니다.
  3. nums 배열 역시 오름차순으로 정렬합니다.
  4. 크기가 n이고 0으로 초기화된 결과 배열 res를 준비합니다.
  5. k를 31부터 0까지 1씩 감소시키며 다음 과정을 반복합니다.
    • 새로운 집합 prefixes를 생성하고 포인터 j를 0으로 초기화합니다.
    • 정렬된 각 쿼리 (i, x, limit)에 대해 다음을 수행합니다.
      • j가 m-1 이하이고 nums[j]가 limit 이하인 동안, nums[j]를 오른쪽으로 k비트 시프트한 값을 prefixes에 추가하고 j를 1 증가시킵니다.
      • prefixes가 비어 있다면 조건을 만족하는 원소가 없다는 의미이므로 res[i] := -1로 설정합니다.
      • 그렇지 않으면 res[i]를 왼쪽으로 1비트 시프트한 뒤, target := res[i] XOR 1을 계산합니다. 만약 (x를 오른쪽으로 k비트 시프트한 값) XOR target이 prefixes에 존재한다면 res[i] := target으로 갱신합니다. 이는 해당 비트를 1로 만들 수 있는지 확인하는 결정적 단계입니다.
  6. 모든 반복이 끝나면 res를 반환합니다.

Python 구현 예제

아래 코드를 통해 위 알고리즘이 실제로 어떻게 구현되는지 살펴보겠습니다.

def solve(nums, queries):
    m, n = len(nums), len(queries)
    queries = sorted(((i, x, limit) for i, (x, limit) in enumerate(queries)), key=lambda x: x[2])
    nums = sorted(nums)
    res = [0] * n
    for k in range(31, -1, -1):
        prefixes = set()
        j = 0
        for i, x, limit in queries:
            while j <= m - 1 and nums[j] <= limit:
                prefixes.add(nums[j] >> k)
                j += 1
            if not prefixes:
                res[i] = -1
            else:
                res[i] <<= 1
                target = res[i] ^ 1
                if (x >> k) ^ target in prefixes:
                    res[i] = target
    return res

nums = [0,1,2,3,4]
queries = [[3,1],[1,3],[5,6]]
print(solve(nums, queries))

입력

[0,1,2,3,4], [[3,1],[1,3],[5,6]]

출력

[3, 3, 7]

복잡도 분석 및 마무리

이 알고리즘은 쿼리 정렬에 O(q log q), 배열 정렬에 O(n log n)이 소요되며, 이후 각 비트 자리(최대 32개)마다 배열과 쿼리를 한 번씩 순회하므로 전체 시간 복잡도는 약 O(n log n + q log q + 32×(n+q))입니다. 모든 원소와 쿼리 쌍을 일일이 비교하는 브루트포스 방식(O(n×q))에 비해 훨씬 효율적이므로, 배열과 쿼리의 크기가 클 때 특히 유용합니다. 비트 연산과 정렬, 그리디 선택을 조합하면 복잡한 XOR 문제도 깔끔하게 해결할 수 있습니다.