3067. 在带权树网络中统计可连接服务器对数目

分糖果

• 难度: medium
• 原始链接: https://leetcode.cn/problems/count-pairs-of-connectable-servers-in-a-weighted-tree-network
• 标签: 深度优先遍历

解法一: 深度优先遍历

go

func countPairsOfConnectableServers(edges [][]int, signalSpeed int) []int {
  n := len(edges) + 1
  grid := make([][][]int, n)
  for _, edge := range edges {
    grid[edge[0]] = append(grid[edge[0]], []int{edge[1], edge[2]})
    grid[edge[1]] = append(grid[edge[1]], []int{edge[0], edge[2]})
  }
  var dfs func(int, int, int) int
  dfs = func(cur, parent, sum int) int {
    count := 0
    if sum%signalSpeed == 0 {
      count++
    }
    for _, node := range grid[cur] {
      if node[0] != parent {
        count += dfs(node[0], cur, sum+node[1])
      }
    }
    return count
  }
  ans := make([]int, n)
  for i := 0; i < n; i++ {
    if len(grid[i]) <= 1 {
      continue
    }
    sum := 0
    for _, node := range grid[i] {
      count := dfs(node[0], i, node[1])
      ans[i] += count * sum
      sum += count
    }
  }
  return ans
}