[leetcode #421] Maximum XOR of Two Numbers in an Array

Problem

Given an integer array nums, return the maximum result of nums[i] XOR nums[j], where 0 <= i <= j < n.

Example 1:

Input: nums = [3,10,5,25,2,8]
Output: 28
Explanation: The maximum result is 5 XOR 25 = 28.

Example 2:

Input: nums = [14,70,53,83,49,91,36,80,92,51,66,70]
Output: 127

Constraints:

・ 1 <= nums.length <= 2 * 10⁵
・ 0 <= nums[i] <= 2³¹ - 1

Idea

아이디어는 떠올리기 쉬운데 구현을 어떻게 해야 할 지가 관건인 문제다.

Brute-force로 전체 수에 대해 XOR 연산을 하면 O(n²)이므로 당연히 Time Limit Exceeded가 뜬다.

XOR 연산을 했을 때 최대값은 가장 큰 자리 수가 1이 나오는 것이다. 그래서 자리 수를 내림차순으로 XOR 연산을 하면서 1이 나오면 그 때마다 값을 계산하고 최종적으로 최대값과 비교를 하는 방식을 사용하면 된다.

이를 구현하기 위해 Tree를 사용했다. (구현을 어떻게 할 지 몰라 Leetcode의 discussion을 참조했다.) 우선 새로운 수가 들어오면 Node에 해당 수의 bit를 노드로 만들어 더한다. 수의 범위가 2³¹-1까지이므로 2³⁰까지 bit를 계산하면된다. 트리에 노드를 더할 때 1이 나오면 오른쪽, 0이 나오면 왼쪽에 더하는 방식으로 트리를 만든다.

트리를 다 만든 뒤에는 현재 만들 수 있는 XOR 값을 계산해 최대값과 비교한다. 비교할 때는 트리를 만들 때와 달리 현재 bit 값이 1일 때는 왼쪽을 탐색하고, 0일 때는 오른쪽을 탐색한다. 이는 xor 연산에서 1이 나오려면 bit 값이 서로 달라야 하기 때문이다. 이렇게 하면 Tree에서 현재값으로 만들 수 있는 최대값을 구할 수 있다.

위와 같은 알고리즘을 사용하면 O(n)에 XOR의 최대값을 구할 수 있다.

Solution

class Solution {
    class Node {
        Node left;
        Node right;
    }

    public int findMaximumXOR(int[] nums) {
        Node root = new Node();
        int res = 0;

        for (int num : nums) {
            insert(root, num);
            res = Math.max(res, search(root, num));
        }

        return res;
    }

    private int search(Node node, int val) {
        int xor = 0;

        for (int i=30; i >= 0; i--) {
            int mask = (1 << i);

            if ((val & mask) != 0) {
                if (node.left != null) {
                    xor = xor | mask;
                    node = node.left;
                } else {
                    node = node.right;
                }
            } else {
                if (node.right != null) {
                    xor = xor | mask;
                    node = node.right;
                } else {
                    node = node.left;
                }
            }
        }

        return xor;
    }

    private void insert(Node node, int val) {
        for (int i=30; i >=0; i--) {
            int mask = (1 << i);

            if ((val & mask) != 0) {
                if (node.right == null) {
                    node.right = new Node();
                }
                node = node.right;
            } else {
                if (node.left == null) {
                    node.left = new Node();
                }
                node = node.left;
            }
        }
    }
}

Reference

https://leetcode.com/problems/maximum-xor-of-two-numbers-in-an-array/