题解 | 在树上游玩

package main

import (
	"fmt"
    "math/big"
)
// 解题思路
// 1，找到所有联通的标记点组
// 2，需要染色的最少量就是组的数量
// 3，可能的方案就是连接到每个组内所有点的普通点的数量相乘
// 4，题目的恶心之处，结果64位都存不下，需要一个超大整数类型来计算
// 5，go语言版本1秒内算不完

type Node struct {
    id int
    init bool
}
func main() {
	var n, k int
	fmt.Scan(&n, &k)
	a := make(map[int]struct{}, k)
	for i := 0; i < k; i++ {
		var v int
		fmt.Scan(&v)
		a[v] = struct{}{}
	}
	link := make(map[int][]int, n)
	for i := 1; i < n; i++ {
		var u, v int
		fmt.Scan(&u, &v)
		link[u] = append(link[u], v)
		link[v] = append(link[v], u)
	}
	stack := make([]*Node, 0, n)
	for u := range link {
		stack = append(stack, &Node{u, true})
	}
	visited := make(map[int]struct{})
	group := make([]map[int]struct{}, 0)
	var g map[int]struct{}
	for len(stack) > 0 {
		l := len(stack)
		// top
		item := stack[l-1]
		// pop
		stack = stack[:l-1]
        u := item.id
		// 已访问点跳过
		if _, ok := visited[u]; ok {
			continue
		}
		// 非标记点跳过
		if _, ok := a[u]; !ok {
            g = nil
			continue
		}
        visited[u] = struct{}{}
		vec := link[u]
		if g == nil || item.init {
			g = make(map[int]struct{})
			group = append(group, g)
		}
        g[u] = struct{}{}
		for _, v := range vec {
            if _, ok := visited[v]; ok {
                continue
            }
			if _, ok := a[v]; !ok {
				continue
			}
			stack = append(stack, &Node{v, false})
		}
	}
    // for k, v := range group {
    //     fmt.Printf("%+v:%+v\n", k, v)
    // }

    cost := len(group)
    cnt := big.NewInt(1)
    for _, v := range group {
        flag := make(map[int]struct{})
        for u := range v {
            for _, id := range link[u] {
                // 非标记点
                if _, ok := a[id]; !ok {
                    flag[id] = struct{}{}
                }
            }
        }
        cnt.Mul(cnt, big.NewInt(int64(len(flag))))
    }
	fmt.Printf("%d %d\n", cost, cnt.Mod(cnt, big.NewInt(1e9+7)))
}