2581. 计算可能的根节点数

Hard

题目描述

Alice 有一个标记为 0 到 n - 1 的 n 个节点的无向树。该树表示为长度为 n - 1 的二维整数数组 edges，其中 edges[i] = [ai, bi] 表示树中节点 ai 和 bi 之间有一条边。

Alice 希望 Bob 找到树的根。她允许 Bob 对她的树进行几次猜测。在一次猜测中，他执行以下操作：

选择两个不同的整数 u 和 v，使得树中存在边 [u, v]。
他告诉 Alice，u 是树中 v 的父节点。

Bob 的猜测由二维整数数组 guesses 表示，其中 guesses[j] = [uj, vj] 表示 Bob 猜测 uj 是 vj 的父节点。

Alice 比较懒，不会回复 Bob 的每个猜测，而只是说他的猜测中至少有 k 个是正确的。

给定二维整数数组 edges、guesses 和整数 k，返回可以作为 Alice 的树的根的可能节点数。如果没有这样的树，则返回 0。

示例 1：

输入：edges = [[0,1],[1,2],[1,3],[4,2]], guesses = [[1,3],[0,1],[1,0],[2,4]], k = 3
输出：3
解释：
Root = 0, 正确猜测 = [1,3], [0,1], [2,4]
Root = 1, 正确猜测 = [1,3], [1,0], [2,4]
Root = 2, 正确猜测 = [1,3], [1,0], [2,4]
Root = 3, 正确猜测 = [1,0], [2,4]
Root = 4, 正确猜测 = [1,3], [1,0]
将 0、1 或 2 作为根节点都会导致 3 个正确猜测。

示例 2：

输入：edges = [[0,1],[1,2],[2,3],[3,4]], guesses = [[1,0],[3,4],[2,1],[3,2]], k = 1
输出：5
解释：
Root = 0, 正确猜测 = [3,4]
Root = 1, 正确猜测 = [1,0], [3,4]
Root = 2, 正确猜测 = [1,0], [2,1], [3,4]
Root = 3, 正确猜测 = [1,0], [2,1], [3,2], [3,4]
Root = 4, 正确猜测 = [1,0], [2,1], [3,2]
将任何节点作为根都至少有 1 个正确猜测。

约束条件：

edges.length == n - 1
2 <= n <= 10^5
1 <= guesses.length <= 10^5
0 <= ai, bi, uj, vj <= n - 1
ai != bi
uj != vj
edges 表示一个有效的树
guesses[j] 是树的一条边
guesses 是唯一的
0 <= k <= guesses.length

解题思路

这是一个典型的换根DP问题。我们需要计算每个节点作为根时正确猜测的数量，然后统计满足条件的根节点数。

核心思路：

建图和预处理：构建无向图，并将猜测存储在哈希集合中便于快速查询。
初始根计算：选择任意节点（如节点0）作为初始根，通过DFS计算以它为根时的正确猜测数量。
换根DP：利用换根的特性，当根从节点u转移到相邻节点v时：
- 如果原来猜测[u,v]是正确的，换根后变为错误，正确数-1
- 如果原来猜测[v,u]是错误的，换根后变为正确，正确数+1
状态转移：对于每个节点，递归地计算其作为根时的正确猜测数，并统计满足条件的节点数。

算法步骤：

首先以节点0为根进行DFS，计算初始正确猜测数
然后从节点0开始进行换根DFS，利用状态转移公式更新每个节点的正确猜测数
统计正确猜测数≥k的节点个数

时间复杂度O(n)，空间复杂度O(n)，这是最优解法。

代码实现

class Solution {
public:
    int rootCount(vector<vector<int>>& edges, vector<vector<int>>& guesses, int k) {
        int n = edges.size() + 1;
        vector<vector<int>> graph(n);
        set<pair<int, int>> guessSet;
        
        for (auto& edge : edges) {
            graph[edge[0]].push_back(edge[1]);
            graph[edge[1]].push_back(edge[0]);
        }
        
        for (auto& guess : guesses) {
            guessSet.insert({guess[0], guess[1]});
        }
        
        int correctGuesses = 0;
        function<void(int, int)> dfs = [&](int node, int parent) {
            for (int child : graph[node]) {
                if (child != parent) {
                    if (guessSet.count({node, child})) {
                        correctGuesses++;
                    }
                    dfs(child, node);
                }
            }
        };
        
        dfs(0, -1);
        
        int result = 0;
        function<void(int, int, int)> reroot = [&](int node, int parent, int correct) {
            if (correct >= k) result++;
            
            for (int child : graph[node]) {
                if (child != parent) {
                    int newCorrect = correct;
                    if (guessSet.count({node, child})) newCorrect--;
                    if (guessSet.count({child, node})) newCorrect++;
                    reroot(child, node, newCorrect);
                }
            }
        };
        
        reroot(0, -1, correctGuesses);
        return result;
    }
};

class Solution:
    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:
        n = len(edges) + 1
        graph = [[] for _ in range(n)]
        guess_set = set()
        
        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)
        
        for u, v in guesses:
            guess_set.add((u, v))
        
        correct_guesses = 0
        
        def dfs(node, parent):
            nonlocal correct_guesses
            for child in graph[node]:
                if child != parent:
                    if (node, child) in guess_set:
                        correct_guesses += 1
                    dfs(child, node)
        
        dfs(0, -1)
        
        result = 0
        
        def reroot(node, parent, correct):
            nonlocal result
            if correct >= k:
                result += 1
            
            for child in graph[node]:
                if child != parent:
                    new_correct = correct
                    if (node, child) in guess_set:
                        new_correct -= 1
                    if (child, node) in guess_set:
                        new_correct += 1
                    reroot(child, node, new_correct)
        
        reroot(0, -1, correct_guesses)
        return result

public class Solution {
    public int RootCount(int[][] edges, int[][] guesses, int k) {
        int n = edges.Length + 1;
        List<int>[] graph = new List<int>[n];
        HashSet<(int, int)> guessSet = new HashSet<(int, int)>();
        
        for (int i = 0; i < n; i++) {
            graph[i] = new List<int>();
        }
        
        foreach (var edge in edges) {
            graph[edge[0]].Add(edge[1]);
            graph[edge[1]].Add(edge[0]);
        }
        
        foreach (var guess in guesses) {
            guessSet.Add((guess[0], guess[1]));
        }
        
        int correctGuesses = 0;
        
        void Dfs(int node, int parent) {
            foreach (int child in graph[node]) {
                if (child != parent) {
                    if (guessSet.Contains((node, child))) {
                        correctGuesses++;
                    }
                    Dfs(child, node);
                }
            }
        }
        
        Dfs(0, -1);
        
        int result = 0;
        
        void Reroot(int node, int parent, int correct) {
            if (correct >= k) result++;
            
            foreach (int child in graph[node]) {
                if (child != parent) {
                    int newCorrect = correct;
                    if (guessSet.Contains((node, child))) newCorrect--;
                    if (guessSet.Contains((child, node))) newCorrect++;
                    Reroot(child, node, newCorrect);
                }
            }
        }
        
        Reroot(0, -1, correctGuesses);
        return result;
    }
}

var rootCount = function(edges, guesses, k) {
    const n = edges.length + 1;
    const graph = Array.from({length: n}, () => []);
    const guessSet = new Set();
    
    for (const [u, v] of edges) {
        graph[u].push(v);
        graph[v].push(u);
    }
    
    for (const [u, v] of guesses) {
        guessSet.add(`${u},${v}`);
    }
    
    let correctGuesses = 0;
    
    function dfs(node, parent) {
        for (const child of graph[node]) {
            if (child !== parent) {
                if (guessSet.has(`${node},${child}`)) {
                    correctGuesses++;
                }
                dfs(child, node);
            }
        }
    }
    
    dfs(0, -1);
    
    let result = 0;
    
    function reroot(node, parent, correct) {
        if (correct >= k) result++;
        
        for (const child of graph[node]) {
            if (child !== parent) {
                let newCorrect = correct;
                if (guessSet.has(`${node},${child}`)) newCorrect--;
                if (guessSet.has(`${child},${node}`)) newCorrect++;
                reroot(child, node, newCorrect);
            }
        }
    }
    
    reroot(0, -1, correctGuesses);
    return result;
};

在 LeetCode 上运行：力扣中文版 LeetCode 国际版

复杂度分析

复杂度类型	值	说明
时间复杂度	O(n + m)	其中 n 是节点数，m 是猜测数。需要遍历所有边和猜测各一次
空间复杂度	O(n + m)	图的邻接表存储需要 O(n)，猜测集合存储需要 O(m)，递归栈深度 O(n)

题目描述#

解题思路#

代码实现#

复杂度分析#

相关题目#

题目描述

解题思路

代码实现

复杂度分析

相关题目