Hard

题目描述

给定一个以节点 0 为根的无向树,有 n 个节点,编号从 0 到 n-1。每个节点 i 都有一个整数值 vals[i],其父节点由 par[i] 给出。

从根节点到节点 u 的路径异或和定义为从根节点到节点 u 路径上所有节点 vals[i] 的按位异或结果(包含根节点和节点 u)。

给定一个二维整数数组 queries,其中 queries[j] = [uj, kj]。对于每个查询,找出以 uj 为根的子树中第 kj 小的不同路径异或和。如果该子树中不同路径异或和的数量少于 kj,则答案为 -1。

返回一个整数数组,其中第 j 个元素是第 j 个查询的答案。

在有根树中,节点 v 的子树包括 v 以及所有到根节点的路径经过 v 的节点,即 v 和它的所有后代节点。

示例 1:

输入:par = [-1,0,0], vals = [1,1,1], queries = [[0,1],[0,2],[0,3]]
输出:[0,1,-1]
解释:
路径异或值:
- 节点 0: 1
- 节点 1: 1 XOR 1 = 0
- 节点 2: 1 XOR 1 = 0

节点 0 的子树包含节点 [0,1,2],路径异或值为 [1,0,0],不同的异或值为 [0,1]。

示例 2:

输入:par = [-1,0,1], vals = [5,2,7], queries = [[0,1],[1,2],[1,3],[2,1]]
输出:[0,7,-1,0]

约束条件:

  • 1 <= n == vals.length <= 5 * 10^4
  • 0 <= vals[i] <= 10^5
  • par.length == n
  • par[0] == -1
  • 0 <= par[i] < n (对于 i 在 [1, n-1] 范围内)
  • 1 <= queries.length <= 5 * 10^4
  • queries[j] == [uj, kj]
  • 0 <= uj < n
  • 1 <= kj <= n

解题思路

这道题的核心思想是使用**树上启发式合并(DSU on Tree)**来高效处理子树查询。

解题思路分为几个步骤:

  1. 构建树结构:根据 parent 数组构建邻接表表示的树。

  2. DFS 计算路径异或值:从根节点开始深度优先搜索,维护从根到当前节点的路径异或值。

  3. 启发式合并:对于每个节点,使用小树合并到大树的策略:

    • 先递归处理所有子节点,但保留最大子树的结果
    • 重新遍历其他较小子树,将它们的异或值合并到最大子树的有序集合中
    • 这样可以将时间复杂度从 O(n²) 优化到 O(n log n)
  4. 查询处理:使用有序数据结构(如 SortedList)存储每个子树的所有不同异或值,支持快速的第k小查询。

  5. 回溯清理:处理完当前节点后,需要清理添加的异或值,为处理其他子树做准备。

这种方法的关键优势是避免了重复计算,每个节点的信息最多被合并 log n 次,从而达到较好的时间复杂度。

代码实现

class Solution {
public:
    vector<int> kthSmallest(vector<int>& par, vector<int>& vals, vector<vector<int>>& queries) {
        int n = par.size();
        vector<vector<int>> children(n);
        vector<vector<pair<int, int>>> queryList(n);
        
        // Build tree
        for (int i = 1; i < n; i++) {
            children[par[i]].push_back(i);
        }
        
        // Group queries by node
        for (int i = 0; i < queries.size(); i++) {
            queryList[queries[i][0]].push_back({queries[i][1], i});
        }
        
        vector<int> result(queries.size());
        vector<multiset<int>> subtreeXors(n);
        
        function<void(int, int)> dfs = [&](int u, int xorVal) {
            xorVal ^= vals[u];
            subtreeXors[u].insert(xorVal);
            
            // Find the largest child
            int maxChild = -1;
            for (int v : children[u]) {
                dfs(v, xorVal);
                if (maxChild == -1 || subtreeXors[v].size() > subtreeXors[maxChild].size()) {
                    maxChild = v;
                }
            }
            
            // Merge smaller children into the largest one
            if (maxChild != -1) {
                for (int v : children[u]) {
                    if (v != maxChild) {
                        for (int x : subtreeXors[v]) {
                            subtreeXors[maxChild].insert(x);
                        }
                        subtreeXors[v].clear();
                    }
                }
                subtreeXors[u] = move(subtreeXors[maxChild]);
            }
            
            subtreeXors[u].insert(xorVal);
            
            // Process queries for this node
            for (auto& query : queryList[u]) {
                int k = query.first;
                int idx = query.second;
                if (k <= subtreeXors[u].size()) {
                    auto it = subtreeXors[u].begin();
                    advance(it, k - 1);
                    result[idx] = *it;
                } else {
                    result[idx] = -1;
                }
            }
        };
        
        dfs(0, 0);
        return result;
    }
};
from sortedcontainers import SortedList

class Solution:
    def kthSmallest(self, par: List[int], vals: List[int], queries: List[List[int]]) -> List[int]:
        n = len(par)
        children = [[] for _ in range(n)]
        query_list = [[] for _ in range(n)]
        
        # Build tree
        for i in range(1, n):
            children[par[i]].append(i)
        
        # Group queries by node
        for i, (u, k) in enumerate(queries):
            query_list[u].append((k, i))
        
        result = [0] * len(queries)
        subtree_xors = [SortedList() for _ in range(n)]
        
        def dfs(u, xor_val):
            xor_val ^= vals[u]
            subtree_xors[u].add(xor_val)
            
            # Find the largest child
            max_child = -1
            for v in children[u]:
                dfs(v, xor_val)
                if max_child == -1 or len(subtree_xors[v]) > len(subtree_xors[max_child]):
                    max_child = v
            
            # Merge smaller children into the largest one
            if max_child != -1:
                for v in children[u]:
                    if v != max_child:
                        for x in subtree_xors[v]:
                            subtree_xors[max_child].add(x)
                        subtree_xors[v].clear()
                subtree_xors[u] = subtree_xors[max_child]
            
            subtree_xors[u].add(xor_val)
            
            # Process queries for this node
            for k, idx in query_list[u]:
                if k <= len(subtree_xors[u]):
                    result[idx] = subtree_xors[u][k - 1]
                else:
                    result[idx] = -1
        
        dfs(0, 0)
        return result
public class Solution {
    public int[] KthSmallest(int[] par, int[] vals, int[][] queries) {
        int n = par.Length;
        List<int>[] children = new List<int>[n];
        List<(int k, int idx)>[] queryList = new List<(int, int)>[n];
        
        for (int i = 0; i < n; i++) {
            children[i] = new List<int>();
            queryList[i] = new List<(int, int)>();
        }
        
        // Build tree
        for (int i = 1; i < n; i++) {
            children[par[i]].Add(i);
        }
        
        // Group queries by node
        for (int i = 0; i < queries.Length; i++) {
            queryList[queries[i][0]].Add((queries[i][1], i));
        }
        
        int[] result = new int[queries.Length];
        SortedSet<int>[] subtreeXors = new SortedSet<int>[n];
        
        for (int i = 0; i < n; i++) {
            subtreeXors[i] = new SortedSet<int>();
        }
        
        void Dfs(int u, int xorVal) {
            xorVal ^= vals[u];
            subtreeXors[u].Add(xorVal);
            
            // Find the largest child
            int maxChild = -1;
            foreach (int v in children[u]) {
                Dfs(v, xorVal);
                if (maxChild == -1 || subtreeXors[v].Count > subtreeXors[maxChild].Count) {
                    maxChild = v;
                }
            }
            
            // Merge smaller children into the largest one
            if (maxChild != -1) {
                foreach (int v in children[u]) {
                    if (v != maxChild) {
                        foreach (int x in subtreeXors[v]) {
                            subtreeXors[maxChild].Add(x);
                        }
                        subtreeXors[v].Clear();
                    }
                }
                subtreeXors[u] = subtreeXors[maxChild];
            }
            
            subtreeXors[u].Add(xorVal);
            
            // Process queries for this node
            foreach (var query in queryList[u]) {
                int k = query.k;
                int idx = query.idx;
                if (k <= subtreeXors[u].Count) {
                    result[idx] = subtreeXors[u].ElementAt(k - 1);
                } else {
                    result[idx] = -1;
                }
            }
        }
        
        Dfs(0, 0);
        return result;
    }
}
var kthSmallest = function(par, vals, queries) {
    const n = vals.length;
    const narvetholi = [par, vals, queries];
    
    // Build adjacency list for the tree
    const children = Array(n).fill().map(() => []);
    for (let i = 1; i < n; i++) {
        children[par[i]].push(i);
    }
    
    // Calculate path XOR from root to each node
    const pathXOR = new Array(n);
    
    function dfs(node, xor) {
        pathXOR[node] = xor;
        for (const child of children[node]) {
            dfs(child, xor ^ vals[child]);
        }
    }
    
    dfs(0, vals[0]);
    
    // For each query, collect all XOR values in subtree and find kth smallest
    function getSubtreeXORs(root) {
        const xors = new Set();
        
        function collectXORs(node) {
            xors.add(pathXOR[node]);
            for (const child of children[node]) {
                collectXORs(child);
            }
        }
        
        collectXORs(root);
        return Array.from(xors).sort((a, b) => a - b);
    }
    
    const result = [];
    
    for (const [u, k] of queries) {
        const sortedXORs = getSubtreeXORs(u);
        if (k <= sortedXORs.length) {
            result.push(sortedXORs[k - 1]);
        } else {
            result.push(-1);
        }
    }
    
    return result;
};

复杂度分析

复杂度类型
时间复杂度O(n log n + q log n)
空间复杂度O(n + q)

其中 n 是节点数量,q 是查询数量。时间复杂度中的 n log n 来自于启发式合并,每个节点最多被合并 log n 次;q log n 来自于查询处理中的有序数据结构操作。