leetcode-go/medium/310.minimum-height-trees.go

func findMinHeightTreesBFS(n int, edges [][]int) (res []int) { // BFS, start from leaf
	if n <= 2 {
		for i := 0; i < n; i++ {
			res = append(res, i)
		}
		return
	}
	adj := make([][]int, n)
	inDegree := make(map[int]int)
	for i := range edges {
		a, b := edges[i][0], edges[i][1]
		adj[a] = append(adj[a], b)
		adj[b] = append(adj[b], a)
		inDegree[a]++
		inDegree[b]++
	}
	for 2 < len(inDegree) {
		li := make([]int, 0)
		for k, v := range inDegree {
			if v == 1 {
				li = append(li, k)
			}
		}
		for _, i := range li {
			for _, j := range adj[i] {
				if v, ok := inDegree[j]; ok {
					inDegree[j] = v - 1
				}
			}
			delete(inDegree, i)
		}
	}
	for k, _ := range inDegree {
		res = append(res, k)
	}
	return
}

type pair struct {
	_1 int
	_2 int
}

func findMinHeightTrees(n int, edges [][]int) (res []int) {
	adj := make([][]int, n)
	for i := range edges {
		a, b := edges[i][0], edges[i][1]
		adj[a] = append(adj[a], b)
		adj[b] = append(adj[b], a)
	}
	minDepth := 1<<32 - 1
	color := make([]int, n)
	m := make(map[pair]int)
	for i := 0; i < n; i++ {
		maxDepth := getDepth(adj, color, i, &m)
		if maxDepth < minDepth {
			minDepth = maxDepth
			res = []int{i}
		} else if maxDepth == minDepth {
			res = append(res, i)
		}
	}
	return
}

func getDepth(adj [][]int, color []int, id int, m *map[pair]int) (depth int) {
	if len(adj[id]) == 1 && color[adj[id][0]] == 1 {
		return
	}
	color[id] = 1
	for _, i := range adj[id] {
		if color[i] != 0 {
			continue
		}
		if d, ok := (*m)[pair{id, i}]; ok {
			depth = maxInt(depth, d)
			continue
		}
		d := getDepth(adj, color, i, m)
		(*m)[pair{id, i}] = d
		depth = maxInt(depth, d)
	}
	color[id] = 0
	return depth + 1
}

func maxInt(x, y int) int {
	if x < y {
		return y
	}
	return x
}