Submission #2960182

---

Source Code Expand

package main

import "fmt"

type Node struct {
	next    map[int]bool
	checked bool
	values  map[int]bool
	value   int
}

var nodes []*Node
var roots = make(map[int]bool)

// 2018-08-04T00:24:54+0800
// 2018-08-04T00:41:59+0800
// 2018-08-04T01:37:57+0800 # Paused for sleeping.
// 2018-08-04T13:05:41+0800 # Continue
// 2018-08-04T13:30:05+0800 # Paused
// 2018-08-05T23:13:36+0800 # Continue
// 2018-08-05T23:49:41+0800
func main() {
	var N, K int
	fmt.Scanf("%d", &N)
	nodes = make([]*Node, N)
	for i := 0; i < N; i++ {
		nodes[i] = &Node{
			make(map[int]bool),
			false,
			make(map[int]bool),
			0,
		}
	}
	a, b := 0, 0
	for i := 0; i < N-1; i++ {
		fmt.Scanf("%d%d", &a, &b)
		nodes[a-1].next[b-1] = true
		nodes[b-1].next[a-1] = true
	}

	// Building the tree.
	//buildTree(0)

	fmt.Scanf("%d", &K)
	for i := 0; i < K; i++ {
		fmt.Scanf("%d%d", &a, &b)
		nodes[a-1].checked = true
		nodes[a-1].values[b] = true
		nodes[a-1].value = b
		roots[a-1] = true
	}
	for i := range roots {
		if !bfs(i, i) {
			fmt.Println("No")
			return
		}
	}
	for i := 0; i < N; i++ {
		if len(nodes[i].values) == 1 {
			roots[i] = true
			for k := range nodes[i].values {
				nodes[i].value = k
			}

		}
	}
	for i := range roots {
		setValues(i, i)
	}
	fmt.Println("Yes")
	for i := 0; i < N; i++ {
		fmt.Println(nodes[i].value)
	}
}

func setValues(i, parent int) {
	for ii := range nodes[i].next {
		if roots[ii] || ii == parent {
			continue
		}
		nodes[ii].value = nodes[i].value + 1
		roots[ii] = true
		setValues(ii, parent)
	}
}

func buildTree(parent int) {
	for i := range nodes[parent].next {
		delete(nodes[i].next, parent)
		buildTree(i)
	}
}

func bfs(root, parent int) bool {
	node := nodes[root]
	for i := range node.next {
		if i == parent {
			// Do not bread search back.
			continue
		}
		pv := potentialValues(node.values)
		if nodes[i].checked {
			pv := merge(pv, nodes[i].values)
			if len(pv) == 0 {
				return false
			}
			nodes[i].values = pv
		} else {
			nodes[i].values = pv
		}
		nodes[i].checked = true
	}
	for i := range node.next {
		if i == parent {
			// Do not bread search back.
			continue
		}
		bfs(i, root)
	}
	return true
}

func potentialValues(pv map[int]bool) map[int]bool {
	m := make(map[int]bool, len(pv)*2)
	for i := range pv {
		m[i-1] = true
		m[i+1] = true
	}
	return m
}

func merge(pv, v2 map[int]bool) map[int]bool {
	m := make(map[int]bool)
	for i := range pv {
		if v2[i] {
			m[i] = true
		}
	}
	return m
}

Submission Info

Judge Result