小红书9-1笔试第三题

package xhs

import (
"fmt"
"math"
)

func max(a, b int) int {
if a > b {
return a
}
return b
}

func dfs(node int, parent int, colors []rune, adjList map[int][]int, blackCount []int) int {
count := 0
if colors[node-1] == 'B' {
count = 1
}

for _, child := range adjList[node] {
if child == parent {
continue
}
count += dfs(child, node, colors, adjList, blackCount)
}

blackCount[node] = count
return count
}

func findMaxBlackNodes(n int, colors string, edges [][]int) int {
adjList := make(map[int][]int)
blackCount := make([]int, n+1)

// 建立邻接表
for _, edge := range edges {
u, v := edge[0], edge[1]
adjList[u] = append(adjList[u], v)
adjList[v] = append(adjList[v], u)
}

// 颜色转换为字符数组
colorRunes := []rune(colors)

// 从根节点开始进行DFS，计算黑色节点数量
dfs(1, -1, colorRunes, adjList, blackCount)

// 找到删除一个红色节点后的最大黑色节点数
maxBlack := math.MinInt32

for i := 1; i <= n; i++ {

for _, node := range adjList[i] {
if colorRunes[node-1] == 'R' {
maxBlack = max(maxBlack, blackCount[i]-blackCount[node])
}
}
}

return maxBlack
}

func XHS9_1_3main() {
n := 10
colors := &quot;RRBBBBBBBB&quot;
edges := [][]int{
{1, 2},
{1, 3},
{1, 4},
{1, 5},
{1, 7},
{2, 8},
{4, 6},
{4, 10},
{6, 9},
}

result := findMaxBlackNodes(n, colors, edges)
fmt.Println(result) // 输出 7
}