algorithms/leetcode/2049-CountNodesWithTheHighestScore.go at master · freewu/algorithms · GitHub

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
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
package main

// 2049. Count Nodes With the Highest Score
// There is a binary tree rooted at 0 consisting of n nodes.
// The nodes are labeled from 0 to n - 1.
// You are given a 0-indexed integer array parents representing the tree, where parents[i] is the parent of node i.
// Since node 0 is the root, parents[0] == -1.

// Each node has a score. To find the score of a node, consider if the node and the edges connected to it were removed.
// The tree would become one or more non-empty subtrees.
// The size of a subtree is the number of the nodes in it.
// The score of the node is the product of the sizes of all those subtrees.

// Return the number of nodes that have the highest score.

// Example 1:
// <img src="https://assets.leetcode.com/uploads/2021/10/03/example-1.png" />
// Input: parents = [-1,2,0,2,0]
// Output: 3
// Explanation:
// - The score of node 0 is: 3 * 1 = 3
// - The score of node 1 is: 4 = 4
// - The score of node 2 is: 1 * 1 * 2 = 2
// - The score of node 3 is: 4 = 4
// - The score of node 4 is: 4 = 4
// The highest score is 4, and three nodes (node 1, node 3, and node 4) have the highest score.

// Example 2:
// <img src="https://assets.leetcode.com/uploads/2021/10/03/example-2.png" />
// Input: parents = [-1,2,0]
// Output: 2
// Explanation:
// - The score of node 0 is: 2 = 2
// - The score of node 1 is: 2 = 2
// - The score of node 2 is: 1 * 1 = 1
// The highest score is 2, and two nodes (node 0 and node 1) have the highest score.

// Constraints:
//     n == parents.length
//     2 <= n <= 10^5
//     parents[0] == -1
//     0 <= parents[i] <= n - 1 for i != 0
//     parents represents a valid binary tree.

import "fmt"

func countHighestScoreNodes(parents []int) int {
    score, count, n := 0, 0, len(parents)
    tree := make([][]int, n)
    for i, p := range parents {
        if p < 0 { continue }
        tree[p] = append(tree[p], i)
    }
    var dfs func(node int) int
    dfs = func(node int)  int { // dfs return tree size of node
        sum, product := 1, 1
        for _, child := range tree[node] {
            v := dfs(child)
            sum += v
            product *= v
        }
        if n - sum > 0 {
            product *= (n - sum)
        }
        if product == score {
            count++
        }
        if product > score {
            score, count = product, 1
        }
        return sum
    }
    dfs(0)
    return count
}

func countHighestScoreNodes1(parents []int) int {
    res, mx, n := 0, -1_000_000_000, len(parents)
    adjList := make([][]int, n)
    for i := 1; i < n; i++ {
        adjList[parents[i]] = append(adjList[parents[i]], i)
    }
    var dfs func(curr, n int) int
    dfs = func(curr, n int) int {
        l, r, score := 0, 0, 1
        if len(adjList[curr]) > 0 { l = dfs(adjList[curr][0], n) }
        if len(adjList[curr]) > 1 { r = dfs(adjList[curr][1], n) }
        if l > 0 { score *= l }
        if r > 0 { score *= r }
        if n - l - r - 1 > 0 { score *= (n - l - r - 1) }
        if score == mx {
            res++
        } else if score > mx {
            mx, res = score, 1
        }
        return 1 + l + r
    }
    dfs(0, n)
    return res
}

func main() {
    // Example 1:
    // <img src="https://assets.leetcode.com/uploads/2021/10/03/example-1.png" />
    // Input: parents = [-1,2,0,2,0]
    // Output: 3
    // Explanation:
    // - The score of node 0 is: 3 * 1 = 3
    // - The score of node 1 is: 4 = 4
    // - The score of node 2 is: 1 * 1 * 2 = 2
    // - The score of node 3 is: 4 = 4
    // - The score of node 4 is: 4 = 4
    // The highest score is 4, and three nodes (node 1, node 3, and node 4) have the highest score.
    fmt.Println(countHighestScoreNodes([]int{-1,2,0,2,0})) // 3
    // Example 2:
    // <img src="https://assets.leetcode.com/uploads/2021/10/03/example-2.png" />
    // Input: parents = [-1,2,0]
    // Output: 2
    // Explanation:
    // - The score of node 0 is: 2 = 2
    // - The score of node 1 is: 2 = 2
    // - The score of node 2 is: 1 * 1 = 1
    // The highest score is 2, and two nodes (node 0 and node 1) have the highest score.
    fmt.Println(countHighestScoreNodes([]int{-1,2,0})) // 2

    fmt.Println(countHighestScoreNodes1([]int{-1,2,0,2,0})) // 3
    fmt.Println(countHighestScoreNodes1([]int{-1,2,0})) // 2
}