結果

問題 No.1451 集団登校
コンテスト
ユーザー 草苺奶昔
提出日時 2023-03-14 00:21:13
言語 Go
(1.23.4)
結果
AC  
実行時間 136 ms / 2,000 ms
コード長 4,613 bytes
コンパイル時間 11,535 ms
コンパイル使用メモリ 223,476 KB
実行使用メモリ 20,704 KB
最終ジャッジ日時 2024-09-18 07:51:21
合計ジャッジ時間 14,759 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge1
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 27
権限があれば一括ダウンロードができます

ソースコード

diff # 

package main

import (
	"bufio"
	"fmt"
	"os"
	"strings"
)

const MOD int = 1e9 + 7

var INV = Pow(2, MOD-2, MOD)

func main() {
	// https://yukicoder.me/problems/no/1451
	// 初始时有n个人
	// 给定m个操作,每次将i和j所在的班级合并,大小相等时随机选取班长,否则选取较大的班级的班长作为新班级的班长
	// 对每个人,问最后成为班长的概率

	in := bufio.NewReader(os.Stdin)
	out := bufio.NewWriter(os.Stdout)
	defer out.Flush()

	var n, m int
	fmt.Fscan(in, &n, &m)
	edges := make([]Edge, 0, m)
	for i := 0; i < m; i++ {
		var u, v int
		fmt.Fscan(in, &u, &v)
		u, v = u-1, v-1
		edges = append(edges, Edge{u, v, 1})
	}

	tree, roots := ProcessOfMergingTree(n, edges)
	subSize := make([]int, len(tree))
	var dfs1 func(int) int
	dfs1 = func(cur int) int {
		if cur < n { // 原始顶点
			subSize[cur] = 1
		}
		for _, e := range tree[cur] {
			subSize[cur] += dfs1(e.to)
		}
		return subSize[cur]
	}
	for _, root := range roots {
		dfs1(root)
	}

	res := make([]int, n)
	var dfs2 func(int, int)
	dfs2 = func(cur, p int) {
		if cur < n { // 原始顶点
			res[cur] = p
			return
		}
		if len(tree[cur]) == 1 { // 只有一个子节点
			dfs2(tree[cur][0].to, p)
			return
		}
		left, right := tree[cur][0].to, tree[cur][1].to // 两个子节点
		if subSize[left] > subSize[right] {
			dfs2(left, p)
		} else if subSize[left] < subSize[right] {
			dfs2(right, p)
		} else {
			dfs2(left, p*INV%MOD)
			dfs2(right, p*INV%MOD)
		}
	}
	for _, root := range roots {
		dfs2(root, 1)
	}

	for _, v := range res {
		fmt.Fprintln(out, v)
	}
}

type Edge struct{ from, to, weight int }

// 表示合并过程的树,按照edges中边的顺序合并顶点.
//  返回: 树的有向图邻接表, 新图中的各个根节点.
func ProcessOfMergingTree(n int, edges []Edge) (tree [][]Edge, roots []int) {
	parent := make([]int, 2*n-1)
	for i := range parent {
		parent[i] = i
	}

	tree = make([][]Edge, 2*n-1)
	uf := NewUnionFindArray(n)
	aux := n
	for _, e := range edges {
		from, to := e.from, e.to
		f := func(big, small int) {
			w, p1, p2 := e.weight, parent[big], parent[small]
			tree[aux] = append(tree[aux], Edge{aux, p1, w})
			tree[aux] = append(tree[aux], Edge{aux, p2, w})
			parent[p1], parent[p2] = aux, aux
			parent[big], parent[small] = aux, aux
			aux++
		}
		uf.UnionWithCallback(from, to, f)
	}

	tree = tree[:aux]
	for i := 0; i < aux; i++ {
		if parent[i] == i {
			roots = append(roots, i)
		}
	}
	return
}

func Pow(base, exp, mod int) int {
	base %= mod
	res := 1
	for ; exp > 0; exp >>= 1 {
		if exp&1 == 1 {
			res = res * base % mod
		}
		base = base * base % mod
	}
	return res
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

// NewUnionFindWithCallback ...
func NewUnionFindArray(n int) *UnionFindArray {
	parent, rank := make([]int, n), make([]int, n)
	for i := 0; i < n; i++ {
		parent[i] = i
		rank[i] = 1
	}

	return &UnionFindArray{
		Part:   n,
		rank:   rank,
		n:      n,
		parent: parent,
	}
}

type UnionFindArray struct {
	// 连通分量的个数
	Part int

	rank   []int
	n      int
	parent []int
}

func (ufa *UnionFindArray) Union(key1, key2 int) bool {
	root1, root2 := ufa.Find(key1), ufa.Find(key2)
	if root1 == root2 {
		return false
	}

	if ufa.rank[root1] > ufa.rank[root2] {
		root1, root2 = root2, root1
	}
	ufa.parent[root1] = root2
	ufa.rank[root2] += ufa.rank[root1]
	ufa.Part--
	return true
}

func (ufa *UnionFindArray) UnionWithCallback(key1, key2 int, cb func(big, small int)) bool {
	root1, root2 := ufa.Find(key1), ufa.Find(key2)
	if root1 == root2 {
		return false
	}
	if ufa.rank[root1] > ufa.rank[root2] {
		root1, root2 = root2, root1
	}
	ufa.parent[root1] = root2
	ufa.rank[root2] += ufa.rank[root1]
	ufa.Part--
	cb(root2, root1)
	return true
}

func (ufa *UnionFindArray) Find(key int) int {
	for ufa.parent[key] != key {
		ufa.parent[key] = ufa.parent[ufa.parent[key]]
		key = ufa.parent[key]
	}
	return key
}

func (ufa *UnionFindArray) IsConnected(key1, key2 int) bool {
	return ufa.Find(key1) == ufa.Find(key2)
}

func (ufa *UnionFindArray) GetGroups() map[int][]int {
	groups := make(map[int][]int)
	for i := 0; i < ufa.n; i++ {
		root := ufa.Find(i)
		groups[root] = append(groups[root], i)
	}
	return groups
}

func (ufa *UnionFindArray) Size(key int) int {
	return ufa.rank[ufa.Find(key)]
}

func (ufa *UnionFindArray) String() string {
	sb := []string{"UnionFindArray:"}
	for root, member := range ufa.GetGroups() {
		cur := fmt.Sprintf("%d: %v", root, member)
		sb = append(sb, cur)
	}
	sb = append(sb, fmt.Sprintf("Part: %d", ufa.Part))
	return strings.Join(sb, "\n")
}
0