美团提前批2020笔试题第五题-黄黑树高

黄黑树高

我将变量都hard code 到代码里面，a代表节点个数，colors代表每个节点的颜色，parent代表节点下标index+2标号的节点的父节点是parent[index]。该方法使用深度遍历，把所有与当前节点有关的节点都做一次深度计算。
题目的要求直接点就是说，以某个节点为当前根，查看他的子树中和他颜色不一样的子节点，把这些子节点到当前根的距离加起来。
当然了更优秀的人肯定是使用动态规划，但是由于记录某个节点的状态还挺难定义，所以我在做题时没做。

package main

import (
	"fmt"
)

func main() {
	for {
		var a int = 10
		// _, err := fmt.Scan(&a)
		// if err != nil {
		// 	return
		// }
		// colors := make([]int, a)
		// for i := range colors {
		// 	fmt.Scan(&colors[i])
		// }
		// parents := make([]int, a-1)
		// for i := range parents {
		// 	fmt.Scan(&parents[i])
		// }
		colors := []int{0, 0, 1, 0, 0, 1, 1, 1, 0, 0}
		parents := []int{1, 2, 3, 4, 4, 5, 7, 6, 9}
		ans := make([]int, a)
		for i := a; i >= 1; i-- {
			ans[i-1] = test(colors, parents, i, colors[i-1], 0)
		}
		for i := 0; i < len(ans)-1; i++ {
			fmt.Print(ans[i], " ")
		}
		fmt.Print(ans[len(ans)-1])

	}
}
func test(colors, relate []int, parent int, color int, deep int) int {
	var sons []int
	for i, v := range relate {
		if v == parent {
			sons = append(sons, i+2)
		}
	}
	sum := 0
	for _, i := range sons {
		if colors[i-1] != color {
			sum += 1 + deep
		}
		sum += test(colors, relate, i, color, deep+1)
	}
	return sum
}

`答案是17 13 10 6 3 3 0 0 0 0``