1938. Maximum Genetic Difference Query - Leetcode Solution

Code Implementation

class TrieNode:
    def __init__(self):
        self.children = {}
        self.count = 0

class Trie:
    def __init__(self):
        self.root = TrieNode()
        self.L = 20  # Bit length, as n & val up to 10^5

    def insert(self, num):
        node = self.root
        for i in reversed(range(self.L)):
            bit = (num >> i) & 1
            if bit not in node.children:
                node.children[bit] = TrieNode()
            node = node.children[bit]
            node.count += 1

    def remove(self, num):
        node = self.root
        for i in reversed(range(self.L)):
            bit = (num >> i) & 1
            node = node.children[bit]
            node.count -= 1

    def maxXor(self, num):
        node = self.root
        xor = 0
        for i in reversed(range(self.L)):
            bit = (num >> i) & 1
            toggled = 1 - bit
            if toggled in node.children and node.children[toggled].count > 0:
                xor |= (1 << i)
                node = node.children[toggled]
            else:
                node = node.children.get(bit, node)
        return xor

class Solution:
    def maxGeneticDifference(self, parents, queries):
        from collections import defaultdict

        n = len(parents)
        tree = defaultdict(list)
        root = -1
        for child, p in enumerate(parents):
            if p == -1:
                root = child
            else:
                tree[p].append(child)

        qmap = defaultdict(list)
        for idx, (node, val) in enumerate(queries):
            qmap[node].append((val, idx))

        ans = [0] * len(queries)
        trie = Trie()

        def dfs(u):
            trie.insert(u)
            for val, idx in qmap[u]:
                ans[idx] = trie.maxXor(val)
            for v in tree[u]:
                dfs(v)
            trie.remove(u)

        dfs(root)
        return ans

class TrieNode {
public:
    TrieNode* children[2];
    int count;
    TrieNode() : count(0) {
        children[0] = children[1] = nullptr;
    }
};

class Trie {
public:
    TrieNode* root;
    int L;
    Trie() {
        root = new TrieNode();
        L = 20;
    }
    void insert(int num) {
        TrieNode* node = root;
        for (int i = L - 1; i >= 0; --i) {
            int bit = (num >> i) & 1;
            if (!node->children[bit])
                node->children[bit] = new TrieNode();
            node = node->children[bit];
            node->count++;
        }
    }
    void remove(int num) {
        TrieNode* node = root;
        for (int i = L - 1; i >= 0; --i) {
            int bit = (num >> i) & 1;
            node = node->children[bit];
            node->count--;
        }
    }
    int maxXor(int num) {
        TrieNode* node = root;
        int xorVal = 0;
        for (int i = L - 1; i >= 0; --i) {
            int bit = (num >> i) & 1;
            int toggled = 1 - bit;
            if (node->children[toggled] && node->children[toggled]->count > 0) {
                xorVal |= (1 << i);
                node = node->children[toggled];
            } else {
                node = node->children[bit];
            }
        }
        return xorVal;
    }
};

class Solution {
public:
    vector maxGeneticDifference(vector& parents, vector>& queries) {
        int n = parents.size();
        vector> tree(n);
        int root = -1;
        for (int i = 0; i < n; ++i) {
            if (parents[i] == -1) root = i;
            else tree[parents[i]].push_back(i);
        }
        unordered_map>> qmap;
        for (int i = 0; i < queries.size(); ++i) {
            qmap[queries[i][0]].emplace_back(queries[i][1], i);
        }
        vector ans(queries.size());
        Trie trie;
        function dfs = [&](int u) {
            trie.insert(u);
            for (auto& q : qmap[u]) {
                ans[q.second] = trie.maxXor(q.first);
            }
            for (int v : tree[u]) dfs(v);
            trie.remove(u);
        };
        dfs(root);
        return ans;
    }
};

class TrieNode {
    TrieNode[] children = new TrieNode[2];
    int count = 0;
}

class Trie {
    TrieNode root = new TrieNode();
    int L = 20;

    void insert(int num) {
        TrieNode node = root;
        for (int i = L - 1; i >= 0; --i) {
            int bit = (num >> i) & 1;
            if (node.children[bit] == null)
                node.children[bit] = new TrieNode();
            node = node.children[bit];
            node.count++;
        }
    }

    void remove(int num) {
        TrieNode node = root;
        for (int i = L - 1; i >= 0; --i) {
            int bit = (num >> i) & 1;
            node = node.children[bit];
            node.count--;
        }
    }

    int maxXor(int num) {
        TrieNode node = root;
        int xorVal = 0;
        for (int i = L - 1; i >= 0; --i) {
            int bit = (num >> i) & 1;
            int toggled = 1 - bit;
            if (node.children[toggled] != null && node.children[toggled].count > 0) {
                xorVal |= (1 << i);
                node = node.children[toggled];
            } else {
                node = node.children[bit];
            }
        }
        return xorVal;
    }
}

class Solution {
    public int[] maxGeneticDifference(int[] parents, int[][] queries) {
        int n = parents.length;
        List[] tree = new ArrayList[n];
        for (int i = 0; i < n; i++) tree[i] = new ArrayList<>();
        int root = -1;
        for (int i = 0; i < n; i++) {
            if (parents[i] == -1) root = i;
            else tree[parents[i]].add(i);
        }
        Map> qmap = new HashMap<>();
        for (int i = 0; i < queries.length; i++) {
            qmap.computeIfAbsent(queries[i][0], k -> new ArrayList<>()).add(new int[]{queries[i][1], i});
        }
        int[] ans = new int[queries.length];
        Trie trie = new Trie();
        dfs(root, tree, qmap, trie, ans);
        return ans;
    }

    void dfs(int u, List[] tree, Map> qmap, Trie trie, int[] ans) {
        trie.insert(u);
        if (qmap.containsKey(u)) {
            for (int[] q : qmap.get(u)) {
                ans[q[1]] = trie.maxXor(q[0]);
            }
        }
        for (int v : tree[u]) dfs(v, tree, qmap, trie, ans);
        trie.remove(u);
    }
}

class TrieNode {
    constructor() {
        this.children = {};
        this.count = 0;
    }
}

class Trie {
    constructor() {
        this.root = new TrieNode();
        this.L = 20;
    }
    insert(num) {
        let node = this.root;
        for (let i = this.L - 1; i >= 0; --i) {
            let bit = (num >> i) & 1;
            if (!(bit in node.children))
                node.children[bit] = new TrieNode();
            node = node.children[bit];
            node.count++;
        }
    }
    remove(num) {
        let node = this.root;
        for (let i = this.L - 1; i >= 0; --i) {
            let bit = (num >> i) & 1;
            node = node.children[bit];
            node.count--;
        }
    }
    maxXor(num) {
        let node = this.root;
        let xor = 0;
        for (let i = this.L - 1; i >= 0; --i) {
            let bit = (num >> i) & 1;
            let toggled = 1 - bit;
            if (toggled in node.children && node.children[toggled].count > 0) {
                xor |= (1 << i);
                node = node.children[toggled];
            } else {
                node = node.children[bit];
            }
        }
        return xor;
    }
}

var maxGeneticDifference = function(parents, queries) {
    const n = parents.length;
    const tree = Array.from({length: n}, () => []);
    let root = -1;
    for (let i = 0; i < n; ++i) {
        if (parents[i] === -1) root = i;
        else tree[parents[i]].push(i);
    }
    const qmap = {};
    for (let i = 0; i < queries.length; ++i) {
        let [node, val] = queries[i];
        if (!(node in qmap)) qmap[node] = [];
        qmap[node].push([val, i]);
    }
    const ans = Array(queries.length);
    const trie = new Trie();

    function dfs(u) {
        trie.insert(u);
        if (qmap[u]) {
            for (const [val, idx] of qmap[u]) {
                ans[idx] = trie.maxXor(val);
            }
        }
        for (const v of tree[u]) dfs(v);
        trie.remove(u);
    }
    dfs(root);
    return ans;
};

Problem Description

Given a rooted tree represented by the array parents (where parents[i] is the parent of node i, and -1 means node i is the root), you are given several queries. Each query is a pair [node, val]. For each query, you must find the maximum value of val XOR x for any ancestor x of node (including node itself as its own ancestor). Return an array where the ith element is the answer to the ith query. Constraints:

Each query must be answered efficiently, considering up to 10⁵ nodes and queries.
Each query's node is a node index in the tree.
Each query's val is an integer value.

Thought Process

The naive way is, for each query, to walk up the tree from node to the root, collecting all ancestors, and for each, compute val XOR ancestor, keeping the maximum. But this is too slow for large trees and many queries, since each query could take O(tree height) time. We need a way to efficiently, for each query, find the ancestor (including the node itself) whose value gives the maximum XOR with val. This is a classic "maximum XOR" problem, but restricted to the set of ancestors on the path from the root to the current node. The key insight is to process queries as we traverse the tree. As we go down the tree, we can keep track of the current path's node values (ancestors) in a data structure that allows fast maximum XOR queries, such as a Trie (prefix tree). For each query at a node, we can ask the Trie for the best XOR with val among all current ancestors.

Solution Approach

To solve the problem efficiently, we use a Trie (prefix tree) and a depth-first search (DFS) traversal of the tree.

Tree Construction: Build the tree from the parents array, storing children for each node.
Query Mapping: For each query, map it to the node it targets. This allows us to process all queries for a node at once during traversal.
Trie for Maximum XOR: Maintain a Trie where each path from root to leaf represents the binary representation of a number (node index). As we traverse down the tree, we insert the current node number into the Trie. When we backtrack (finish processing a subtree), we remove the node from the Trie.
DFS Traversal: Start at the root. For each node:
- Insert the node's value into the Trie.
- For each query at this node, use the Trie to find the ancestor on the current path that gives the maximum val XOR ancestor.
- Recursively process all children.
- After processing children, remove the node's value from the Trie (to keep the path correct).
Result Collection: Store answers in an array, using the query's index to place it in the correct position.

The Trie allows us to find the best XOR in O(log(max value)) time for each query, and the DFS ensures we only keep ancestors on the current path.

Example Walkthrough

Consider:

parents = [-1, 0, 0, 1, 1, 2] (root is 0)
queries = [[3, 5], [5, 2], [4, 6]]

Step-by-step:

Build the tree:
- 0 is root; children: 1, 2
- 1's children: 3, 4
- 2's child: 5
Map queries:
- Node 3: query [3, 5]
- Node 4: query [4, 6]
- Node 5: query [5, 2]
DFS Traversal:
- Start at 0, insert 0 into Trie.
- Go to 1, insert 1 into Trie.
- Go to 3, insert 3 into Trie.
- At 3, process query [3, 5]: Ancestors are [0, 1, 3]. Find max of 5^0=5, 5^1=4, 5^3=6. Max is 6.
- Backtrack: remove 3.
- At 1, go to 4, insert 4. Process [4, 6]: Ancestors [0, 1, 4]. 6^0=6, 6^1=7, 6^4=2. Max is 7.
- Backtrack: remove 4.
- Backtrack: remove 1.
- At 0, go to 2, insert 2.
- Go to 5, insert 5. Process [5, 2]: Ancestors [0, 2, 5]. 2^0=2, 2^2=0, 2^5=7. Max is 7.
- Backtrack: remove 5, then 2, finally 0.
Results:
- Query [3, 5]: 6
- Query [5, 2]: 7
- Query [4, 6]: 7

Time and Space Complexity

Brute-force:

For each query, walk up the tree (O(tree height)), and for each ancestor, compute XOR. O(Q * H), where Q = number of queries, H = tree height. Too slow for large trees.

Optimized (Trie + DFS):

Each node is inserted/removed from the Trie once: O(N * L), where N = number of nodes, L = bit length (about 20 for 10⁵).
Each query does a Trie search: O(Q * L).
Total: O((N + Q) * L).
Space: O(N * L) for the Trie (since at most N nodes in the Trie at any time), plus O(N) for the tree, O(Q) for answers.

Summary

This problem asks for the maximum XOR between a value and any ancestor of a node in a tree, for many queries. The elegant solution uses a Trie to efficiently track all ancestors on the current path during a DFS traversal, enabling fast maximum XOR queries for each node. This approach reduces the time from potentially O(Q * N) to O((N + Q) * L), making it feasible for large inputs. The key insight is to combine tree traversal with a dynamic data structure (the Trie) to answer path-restricted maximum XOR queries efficiently.

AlgoMap.io

Want Help Cracking FAANG?