Skip to content

2581. Count Number of Possible Root Nodes

SourceBiweekly Contest 99 Q4DifficultyHardRating2228

Description

Alice has an undirected tree with n nodes labeled from 0 to n - 1. The tree is represented as a 2D integer array edges of length n - 1 where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree.

Alice wants Bob to find the root of the tree. She allows Bob to make several guesses about her tree. In one guess, he does the following:

  • Chooses two distinct integers u and v such that there exists an edge [u, v] in the tree.
  • He tells Alice that u is the parent of v in the tree.

Bob's guesses are represented by a 2D integer array guesses where guesses[j] = [uj, vj] indicates Bob guessed uj to be the parent of vj.

Alice being lazy, does not reply to each of Bob's guesses, but just says that at least k of his guesses are true.

Given the 2D integer arrays edges, guesses and the integer k, return the number of possible nodes that can be the root of Alice's tree. If there is no such tree, return 0.

 

Example 1:

Input: edges = [[0,1],[1,2],[1,3],[4,2]], guesses = [[1,3],[0,1],[1,0],[2,4]], k = 3
Output: 3
Explanation: 
Root = 0, correct guesses = [1,3], [0,1], [2,4]
Root = 1, correct guesses = [1,3], [1,0], [2,4]
Root = 2, correct guesses = [1,3], [1,0], [2,4]
Root = 3, correct guesses = [1,0], [2,4]
Root = 4, correct guesses = [1,3], [1,0]
Considering 0, 1, or 2 as root node leads to 3 correct guesses.

Example 2:

Input: edges = [[0,1],[1,2],[2,3],[3,4]], guesses = [[1,0],[3,4],[2,1],[3,2]], k = 1
Output: 5
Explanation: 
Root = 0, correct guesses = [3,4]
Root = 1, correct guesses = [1,0], [3,4]
Root = 2, correct guesses = [1,0], [2,1], [3,4]
Root = 3, correct guesses = [1,0], [2,1], [3,2], [3,4]
Root = 4, correct guesses = [1,0], [2,1], [3,2]
Considering any node as root will give at least 1 correct guess.

 

Constraints:

  • edges.length == n - 1
  • 2 <= n <= 105
  • 1 <= guesses.length <= 105
  • 0 <= ai, bi, uj, vj <= n - 1
  • ai != bi
  • uj != vj
  • edges represents a valid tree.
  • guesses[j] is an edge of the tree.
  • guesses is unique.
  • 0 <= k <= guesses.length

Solutions

Solution 1: Tree DP (change root)

Thinking

The root is unknown; guesses are directed edges. A node is valid if, when it is the root, at least \(k\) guesses match. A DFS from every root is \(O(n^2)\).

First count matches \(\textit{cnt}\) with root \(0\). Moving the root to a neighbour \(j\) drops edge \((i,j)\) as a parent-to-child and adds \((j,i)\) if guessed. A second DFS visits every root with this \(O(1)\) update.

First, we traverse the given edge set \(\textit{edges}\) and convert it to an adjacency list \(g\) where \(g[i]\) represents the adjacent nodes of node \(i\). Use a hash map \(gs\) to record the given guess set \(\textit{guesses}\).

Then, we start from node \(0\) and perform a DFS to count the number of edges in \(\textit{guesses}\) among all the nodes that can be reached from node \(0\). We use the variable \(\textit{cnt}\) to record this number.

Next, we start from node \(0\) and perform a DFS to count the number of edges in \(\textit{guesses}\) in each tree with \(0\) as the root. If the number is greater than or equal to \(k\), it means that this node is a possible root node, and we add \(1\) to the answer.

Therefore, the problem becomes to count the number of edges in \(\textit{guesses}\) in each tree with each node as the root. We already know that there are \(\textit{cnt}\) edges in \(\textit{guesses}\) among all the nodes that can be reached from node \(0\). We can maintain this value by adding or subtracting the current edge in \(\textit{guesses}\) in DFS.

Assume that we are currently traversing node \(i\) and \(\textit{cnt}\) represents the number of edges in \(\textit{guesses}\) with \(i\) as the root node. Then, for each adjacent node \(j\) of \(i\), we need to calculate the number of edges in \(\textit{guesses}\) with \(j\) as the root node. If \((i, j)\) is in \(\textit{guesses}\), then there is no edge \((i, j)\) in the tree with \(j\) as the root node, so \(\textit{cnt}\) should decrease by \(1\). If \((j, i)\) is in \(\textit{guesses}\), then there is an extra edge \((i, j)\) in the tree with \(j\) as the root node, so \(\textit{cnt}\) should increase by \(1\). That is, \(f[j] = f[i] + (j, i) \in \textit{guesses} - (i, j) \in \textit{guesses}\). Where \(f[i]\) represents the number of edges in \(\textit{guesses}\) with \(i\) as the root node.

The time complexity is \(O(n + m)\), and the space complexity is \(O(n + m)\), where \(n\) and \(m\) are the lengths of \(\textit{edges}\) and \(\textit{guesses}\) respectively.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
class Solution:
    def rootCount(
        self, edges: List[List[int]], guesses: List[List[int]], k: int
    ) -> int:
        def dfs1(i, fa):
            nonlocal cnt
            for j in g[i]:
                if j != fa:
                    cnt += gs[(i, j)]
                    dfs1(j, i)

        def dfs2(i, fa):
            nonlocal ans, cnt
            ans += cnt >= k
            for j in g[i]:
                if j != fa:
                    cnt -= gs[(i, j)]
                    cnt += gs[(j, i)]
                    dfs2(j, i)
                    cnt -= gs[(j, i)]
                    cnt += gs[(i, j)]

        g = defaultdict(list)
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        gs = Counter((u, v) for u, v in guesses)
        cnt = 0
        dfs1(0, -1)
        ans = 0
        dfs2(0, -1)
        return ans

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
class Solution {
    private List<Integer>[] g;
    private Map<Long, Integer> gs = new HashMap<>();
    private int ans;
    private int k;
    private int cnt;
    private int n;

    public int rootCount(int[][] edges, int[][] guesses, int k) {
        this.k = k;
        n = edges.length + 1;
        g = new List[n];
        Arrays.setAll(g, e -> new ArrayList<>());
        for (var e : edges) {
            int a = e[0], b = e[1];
            g[a].add(b);
            g[b].add(a);
        }
        for (var e : guesses) {
            int a = e[0], b = e[1];
            gs.merge(f(a, b), 1, Integer::sum);
        }
        dfs1(0, -1);
        dfs2(0, -1);
        return ans;
    }

    private void dfs1(int i, int fa) {
        for (int j : g[i]) {
            if (j != fa) {
                cnt += gs.getOrDefault(f(i, j), 0);
                dfs1(j, i);
            }
        }
    }

    private void dfs2(int i, int fa) {
        ans += cnt >= k ? 1 : 0;
        for (int j : g[i]) {
            if (j != fa) {
                int a = gs.getOrDefault(f(i, j), 0);
                int b = gs.getOrDefault(f(j, i), 0);
                cnt -= a;
                cnt += b;
                dfs2(j, i);
                cnt -= b;
                cnt += a;
            }
        }
    }

    private long f(int i, int j) {
        return 1L * i * n + j;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
class Solution {
public:
    int rootCount(vector<vector<int>>& edges, vector<vector<int>>& guesses, int k) {
        int n = edges.size() + 1;
        vector<int> g[n];
        unordered_map<long long, int> gs;
        auto f = [&](int i, int j) {
            return 1LL * i * n + j;
        };
        for (auto& e : edges) {
            int a = e[0], b = e[1];
            g[a].push_back(b);
            g[b].push_back(a);
        }
        for (auto& e : guesses) {
            int a = e[0], b = e[1];
            gs[f(a, b)]++;
        }
        int ans = 0;
        int cnt = 0;

        function<void(int, int)> dfs1 = [&](int i, int fa) {
            for (int& j : g[i]) {
                if (j != fa) {
                    cnt += gs[f(i, j)];
                    dfs1(j, i);
                }
            }
        };

        function<void(int, int)> dfs2 = [&](int i, int fa) {
            ans += cnt >= k;
            for (int& j : g[i]) {
                if (j != fa) {
                    int a = gs[f(i, j)];
                    int b = gs[f(j, i)];
                    cnt -= a;
                    cnt += b;
                    dfs2(j, i);
                    cnt -= b;
                    cnt += a;
                }
            }
        };
        dfs1(0, -1);
        dfs2(0, -1);
        return ans;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
func rootCount(edges [][]int, guesses [][]int, k int) (ans int) {
    n := len(edges) + 1
    g := make([][]int, n)
    gs := map[int]int{}
    for _, e := range edges {
        a, b := e[0], e[1]
        g[a] = append(g[a], b)
        g[b] = append(g[b], a)
    }
    f := func(i, j int) int {
        return i*n + j
    }
    for _, e := range guesses {
        a, b := e[0], e[1]
        gs[f(a, b)]++
    }

    cnt := 0
    var dfs1 func(i, fa int)
    var dfs2 func(i, fa int)
    dfs1 = func(i, fa int) {
        for _, j := range g[i] {
            if j != fa {
                cnt += gs[f(i, j)]
                dfs1(j, i)
            }
        }
    }
    dfs2 = func(i, fa int) {
        if cnt >= k {
            ans++
        }
        for _, j := range g[i] {
            if j != fa {
                a, b := gs[f(i, j)], gs[f(j, i)]
                cnt -= a
                cnt += b
                dfs2(j, i)
                cnt -= b
                cnt += a
            }
        }
    }
    dfs1(0, -1)
    dfs2(0, -1)
    return
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
function rootCount(edges: number[][], guesses: number[][], k: number): number {
    const n = edges.length + 1;
    const g: number[][] = Array.from({ length: n }, () => []);
    const gs: Map<number, number> = new Map();
    const f = (i: number, j: number) => i * n + j;
    for (const [a, b] of edges) {
        g[a].push(b);
        g[b].push(a);
    }
    for (const [a, b] of guesses) {
        const x = f(a, b);
        gs.set(x, gs.has(x) ? gs.get(x)! + 1 : 1);
    }
    let ans = 0;
    let cnt = 0;
    const dfs1 = (i: number, fa: number): void => {
        for (const j of g[i]) {
            if (j !== fa) {
                cnt += gs.get(f(i, j)) || 0;
                dfs1(j, i);
            }
        }
    };
    const dfs2 = (i: number, fa: number): void => {
        ans += cnt >= k ? 1 : 0;
        for (const j of g[i]) {
            if (j !== fa) {
                const a = gs.get(f(i, j)) || 0;
                const b = gs.get(f(j, i)) || 0;
                cnt -= a;
                cnt += b;
                dfs2(j, i);
                cnt -= b;
                cnt += a;
            }
        }
    };
    dfs1(0, -1);
    dfs2(0, -1);
    return ans;
}

Solution 2: Re-rooting DP + Explicit Stack

Thinking

The root is unknown, and each guess is a directed edge. A node is valid when at least \(k\) guesses match after it becomes the root. With \(n\le 10^5\), a separate walk from every root is quadratic, and recursion along a chain overflows Python once the chain reaches length \(1000\). Moving the root across one edge changes only that edge: from parent \(i\) to child \(j\), \((i,j)\) is no longer parent-to-child and \((j,i)\) becomes parent-to-child. The first stack therefore starts at \(0\) and adds each guess whose direction matches the parent-child edge, producing \(\textit{cnt}\). The second stack stores the adjusted count in the frame. A popped node is counted when that value is at least \(k\), and the child receives \(\textit{cnt}-(i,j)+(j,i)\).

Turn \(\textit{edges}\) into an adjacency list \(g\), and store \(\textit{guesses}\) in a hash map \(gs\). An explicit stack walks the tree from node \(0\). When the walk moves to a child \(j\), add the multiplicity of \((i,j)\) in \(gs\) to \(\textit{cnt}\). After this walk, \(\textit{cnt}\) is the number of correct guesses with root \(0\).

The second stack reroots the tree. Each frame stores a node, its parent, and the number of correct guesses when that node is the root. It starts with \((0,-1,\textit{cnt})\). When node \(i\) is popped, add \(1\) to the answer if the stored count is at least \(k\). For a neighbor \(j\), \((i,j)\) is no longer parent-to-child and \((j,i)\) becomes parent-to-child, so the count passed to \(j\) is \(c-gs[(i,j)]+gs[(j,i)]\).

The time complexity is \(O(n + m)\), and the space complexity is \(O(n + m)\), where \(n\) and \(m\) are the lengths of \(\textit{edges}\) and \(\textit{guesses}\) respectively.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class Solution:
    def rootCount(
        self, edges: List[List[int]], guesses: List[List[int]], k: int
    ) -> int:
        g = defaultdict(list)
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        gs = Counter((u, v) for u, v in guesses)
        cnt = 0
        stk = [(0, -1)]
        while stk:
            i, fa = stk.pop()
            for j in g[i]:
                if j != fa:
                    cnt += gs[(i, j)]
                    stk.append((j, i))
        ans = 0
        walk = [(0, -1, cnt)]
        while walk:
            i, fa, c = walk.pop()
            ans += c >= k
            for j in g[i]:
                if j != fa:
                    walk.append((j, i, c - gs[(i, j)] + gs[(j, i)]))
        return ans

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
class Solution {
    public int rootCount(int[][] edges, int[][] guesses, int k) {
        int n = edges.length + 1;
        List<Integer>[] g = new List[n];
        Arrays.setAll(g, e -> new ArrayList<>());
        for (var e : edges) {
            int a = e[0], b = e[1];
            g[a].add(b);
            g[b].add(a);
        }
        Map<Long, Integer> gs = new HashMap<>();
        for (var e : guesses) {
            gs.merge(1L * e[0] * n + e[1], 1, Integer::sum);
        }
        int cnt = 0;
        Deque<int[]> stk = new ArrayDeque<>();
        stk.push(new int[] {0, -1});
        while (!stk.isEmpty()) {
            int[] cur = stk.pop();
            int i = cur[0], fa = cur[1];
            for (int j : g[i]) {
                if (j != fa) {
                    cnt += gs.getOrDefault(1L * i * n + j, 0);
                    stk.push(new int[] {j, i});
                }
            }
        }
        int ans = 0;
        Deque<int[]> walk = new ArrayDeque<>();
        walk.push(new int[] {0, -1, cnt});
        while (!walk.isEmpty()) {
            int[] cur = walk.pop();
            int i = cur[0], fa = cur[1], c = cur[2];
            if (c >= k) {
                ++ans;
            }
            for (int j : g[i]) {
                if (j != fa) {
                    int a = gs.getOrDefault(1L * i * n + j, 0);
                    int b = gs.getOrDefault(1L * j * n + i, 0);
                    walk.push(new int[] {j, i, c - a + b});
                }
            }
        }
        return ans;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
class Solution {
public:
    int rootCount(vector<vector<int>>& edges, vector<vector<int>>& guesses, int k) {
        int n = edges.size() + 1;
        vector<vector<int>> g(n);
        for (auto& e : edges) {
            int a = e[0], b = e[1];
            g[a].push_back(b);
            g[b].push_back(a);
        }
        unordered_map<long long, int> gs;
        for (auto& e : guesses) {
            gs[1LL * e[0] * n + e[1]]++;
        }
        auto get = [&](int i, int j) {
            auto it = gs.find(1LL * i * n + j);
            return it == gs.end() ? 0 : it->second;
        };
        int cnt = 0;
        vector<pair<int, int>> stk{{0, -1}};
        while (!stk.empty()) {
            auto [i, fa] = stk.back();
            stk.pop_back();
            for (int j : g[i]) {
                if (j != fa) {
                    cnt += get(i, j);
                    stk.push_back({j, i});
                }
            }
        }
        int ans = 0;
        vector<array<int, 3>> walk{{0, -1, cnt}};
        while (!walk.empty()) {
            auto [i, fa, c] = walk.back();
            walk.pop_back();
            ans += c >= k;
            for (int j : g[i]) {
                if (j != fa) {
                    walk.push_back({j, i, c - get(i, j) + get(j, i)});
                }
            }
        }
        return ans;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
func rootCount(edges [][]int, guesses [][]int, k int) (ans int) {
    n := len(edges) + 1
    g := make([][]int, n)
    for _, e := range edges {
        a, b := e[0], e[1]
        g[a] = append(g[a], b)
        g[b] = append(g[b], a)
    }
    gs := map[int]int{}
    f := func(i, j int) int {
        return i*n + j
    }
    for _, e := range guesses {
        a, b := e[0], e[1]
        gs[f(a, b)]++
    }
    cnt := 0
    type frame struct{ i, fa int }
    stk := []frame{{0, -1}}
    for len(stk) > 0 {
        cur := stk[len(stk)-1]
        stk = stk[:len(stk)-1]
        for _, j := range g[cur.i] {
            if j != cur.fa {
                cnt += gs[f(cur.i, j)]
                stk = append(stk, frame{j, cur.i})
            }
        }
    }
    type step struct{ i, fa, c int }
    walk := []step{{0, -1, cnt}}
    for len(walk) > 0 {
        cur := walk[len(walk)-1]
        walk = walk[:len(walk)-1]
        if cur.c >= k {
            ans++
        }
        for _, j := range g[cur.i] {
            if j != cur.fa {
                walk = append(walk, step{j, cur.i, cur.c - gs[f(cur.i, j)] + gs[f(j, cur.i)]})
            }
        }
    }
    return
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
function rootCount(edges: number[][], guesses: number[][], k: number): number {
    const n = edges.length + 1;
    const g: number[][] = Array.from({ length: n }, () => []);
    const gs: Map<number, number> = new Map();
    const f = (i: number, j: number) => i * n + j;
    for (const [a, b] of edges) {
        g[a].push(b);
        g[b].push(a);
    }
    for (const [a, b] of guesses) {
        const x = f(a, b);
        gs.set(x, (gs.get(x) || 0) + 1);
    }
    let cnt = 0;
    const stk: number[][] = [[0, -1]];
    while (stk.length) {
        const [i, fa] = stk.pop()!;
        for (const j of g[i]) {
            if (j !== fa) {
                cnt += gs.get(f(i, j)) || 0;
                stk.push([j, i]);
            }
        }
    }
    let ans = 0;
    const walk: number[][] = [[0, -1, cnt]];
    while (walk.length) {
        const [i, fa, c] = walk.pop()!;
        if (c >= k) {
            ans++;
        }
        for (const j of g[i]) {
            if (j !== fa) {
                const a = gs.get(f(i, j)) || 0;
                const b = gs.get(f(j, i)) || 0;
                walk.push([j, i, c - a + b]);
            }
        }
    }
    return ans;
}

Comments