#849475 (Go) No.1976 Cut then Connect

提出ソース

結果

問題	No.1976 Cut then Connect
ユーザー	草苺奶昔
提出日時	2023-03-16 00:18:49
言語	Go (1.22.1)
結果	AC
実行時間	160 ms / 2,000 ms
コード長	7,837 bytes
コンパイル時間	10,114 ms
コンパイル使用メモリ	233,688 KB
実行使用メモリ	32,664 KB
最終ジャッジ日時	2024-11-24 12:19:10
合計ジャッジ時間	13,079 ms
ジャッジサーバーID （参考情報）	judge5 / judge3

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	1 ms 6,820 KB
testcase_01	AC	1 ms 6,820 KB
testcase_02	AC	142 ms 31,808 KB
testcase_03	AC	76 ms 18,936 KB
testcase_04	AC	28 ms 10,012 KB
testcase_05	AC	68 ms 16,432 KB
testcase_06	AC	118 ms 25,004 KB
testcase_07	AC	75 ms 17,020 KB
testcase_08	AC	156 ms 31,004 KB
testcase_09	AC	140 ms 28,816 KB
testcase_10	AC	37 ms 10,112 KB
testcase_11	AC	148 ms 31,092 KB
testcase_12	AC	5 ms 6,816 KB
testcase_13	AC	110 ms 25,008 KB
testcase_14	AC	109 ms 23,360 KB
testcase_15	AC	96 ms 22,156 KB
testcase_16	AC	146 ms 29,956 KB
testcase_17	AC	45 ms 14,272 KB
testcase_18	AC	9 ms 6,820 KB
testcase_19	AC	110 ms 24,880 KB
testcase_20	AC	160 ms 32,664 KB
testcase_21	AC	81 ms 20,304 KB
testcase_22	AC	1 ms 6,816 KB
testcase_23	AC	1 ms 6,816 KB
testcase_24	AC	1 ms 6,820 KB
testcase_25	AC	1 ms 6,820 KB
testcase_26	AC	1 ms 6,816 KB
testcase_27	AC	1 ms 6,816 KB
testcase_28	AC	1 ms 6,816 KB
testcase_29	AC	1 ms 6,816 KB
testcase_30	AC	1 ms 6,816 KB
testcase_31	AC	1 ms 6,820 KB
testcase_32	AC	1 ms 6,820 KB

権限があれば一括ダウンロードができます

ソースコード

raw source code

package main

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

func main() {
	// https://yukicoder.me/problems/no/1976
	// No.1976 Cut then Connect-连边后树的最小直径

	// 给定一棵树, 你可以做以下操作:
	// 从树的无向边中删除一条边, 使得图的连通分量数变为2
	// 对于图的每个连通分量, 选择一个点, 用无向边将这两个点连接起来
	// !问: 操作后的树的直径的最小值是多少?

	// 解:
	// 1. 固定每条删除的边时,求出每个连通分量的直径后可以求出新直径的最小值
	//  !此时答案为 max(x,y,ceil(x/2)+ceil(y/2)+1) 其中x,y为连通分量的直径
	// 2. 对不同的边,考虑换根dp即可

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

	var n int
	fmt.Fscan(in, &n)

	R := NewRerootingEdge(n)
	edges := make([][2]int, n-1)
	for i := 0; i < n-1; i++ {
		var a, b int
		fmt.Fscan(in, &a, &b)
		a, b = a-1, b-1
		edges[i] = [2]int{a, b}
		R.AddEdge(a, b, 1)
	}

	// !E: (每个点处的直径,每个点到最远点的距离)
	R.ReRooting(
		func() E { return E{0, 0} },
		func(dp1, dp2 E) E {
			return E{max(max(dp1.dia, dp2.dia), dp1.dist+dp2.dist), max(dp1.dist, dp2.dist)}
		},
		func(dp E, from int) E { return E{dp.dia, dp.dist + 1} },
		func(dp E, edge Edge) E {
			return dp
		},
	)

	res := n
	for _, e := range edges {
		u, v := e[0], e[1]
		dia1, dia2 := R.Get(u, v).dia, R.Get(v, u).dia
		res = min(res, max(max(dia1, dia2), (dia1+1)/2+(dia2+1)/2+1))
	}
	fmt.Fprintln(out, res)
}

type E = struct{ dia, dist int }
type ReRootingEdge struct {
	tree *_T
	dp1  []E // 边 parent-root 的子树 root 的 dp值
	dp2  []E // 边 parent-root 的子树 parent 的 dp值
	dp   []E // 顶点 v 的子树的 dp值
}

type Edge = struct{ from, to, weight, eid int }

func NewRerootingEdge(n int) *ReRootingEdge {
	return &ReRootingEdge{tree: _NT(n)}
}

func (rr *ReRootingEdge) AddEdge(from, to, weight int) {
	rr.tree.AddEdge(from, to, weight)
}

// !root 作为根节点时, 子树 v 的 dp 值
func (rr *ReRootingEdge) Get(root, v int) E {
	if root == v {
		return rr.dp[v]
	}
	if !rr.tree.IsInSubtree(root, v) {
		return rr.dp1[v]
	}
	w := rr.tree.Jump(v, root, 1)
	return rr.dp2[w]
}

func (rr *ReRootingEdge) ReRooting(
	e func() E,
	op func(dp1, dp2 E) E,
	composition func(dp E, from int) E,
	compositionEdge func(dp E, edge Edge) E,
) []E {
	rr.tree.Build(-1)
	unit := e()
	N := len(rr.tree.Tree)
	dp1, dp2, dp := make([]E, N), make([]E, N), make([]E, N)
	for i := 0; i < N; i++ {
		dp1[i] = unit
		dp2[i] = unit
		dp[i] = unit
	}

	V := rr.tree.idToNode
	par := rr.tree.Parent
	for i := N - 1; i >= 0; i-- {
		v := V[i]
		ch := rr.tree.CollectChild(v)
		n := len(ch)
		x1, x2 := make([]E, n+1), make([]E, n+1)
		for i := range x1 {
			x1[i] = unit
			x2[i] = unit
		}
		for i := 0; i < n; i++ {
			x1[i+1] = op(x1[i], dp2[ch[i]])
		}
		for i := n - 1; i >= 0; i-- {
			x2[i] = op(dp2[ch[i]], x2[i+1])
		}
		for i := 0; i < n; i++ {
			dp2[ch[i]] = op(x1[i], x2[i+1])
		}
		dp[v] = x2[0]
		dp1[v] = composition(dp[v], v)
		for _, e := range rr.tree.Tree[v] {
			to := e.to
			if to == par[v] {
				dp2[v] = compositionEdge(dp1[v], e)
			}
		}
	}

	v := V[0]
	dp[v] = composition(dp[v], v)
	for _, e := range rr.tree.Tree[v] {
		to := e.to
		dp2[to] = composition(dp2[to], v)
	}

	for i := 0; i < N; i++ {
		v := V[i]
		for _, e := range rr.tree.Tree[v] {
			if e.to == par[v] {
				continue
			}
			x := compositionEdge(dp2[e.to], e)
			for _, f := range rr.tree.Tree[e.to] {
				if f.to == par[e.to] {
					continue
				}
				dp2[f.to] = op(dp2[f.to], x)
				dp2[f.to] = composition(dp2[f.to], e.to)
			}
			x = op(dp[e.to], x)
			dp[e.to] = composition(x, e.to)
		}
	}

	rr.dp1, rr.dp2, rr.dp = dp1, dp2, dp
	return dp
}

type _T struct {
	Tree          [][]Edge
	Depth         []int
	Parent        []int
	LID, RID      []int // 欧拉序[in,out)
	idToNode      []int
	top, heavySon []int
	timer         int
	eid           int
}

func _NT(n int) *_T {
	tree := make([][]Edge, n)
	lid := make([]int, n)
	rid := make([]int, n)
	idToNode := make([]int, n)
	top := make([]int, n)      // 所处轻/重链的顶点（深度最小），轻链的顶点为自身
	depth := make([]int, n)    // 深度
	parent := make([]int, n)   // 父结点
	heavySon := make([]int, n) // 重儿子
	for i := range parent {
		parent[i] = -1
	}

	return &_T{
		Tree:     tree,
		Depth:    depth,
		Parent:   parent,
		LID:      lid,
		RID:      rid,
		idToNode: idToNode,
		top:      top,
		heavySon: heavySon,
	}
}

// 添加无向边 u-v, 边权为w.
func (tree *_T) AddEdge(u, v, w int) {
	tree.Tree[u] = append(tree.Tree[u], Edge{from: u, to: v, weight: w, eid: tree.eid})
	tree.Tree[v] = append(tree.Tree[v], Edge{from: v, to: u, weight: w, eid: tree.eid})
	tree.eid++
}

// 添加有向边 u->v, 边权为w.
func (tree *_T) AddDirectedEdge(u, v, w int) {
	tree.Tree[u] = append(tree.Tree[u], Edge{from: u, to: v, weight: w, eid: tree.eid})
	tree.eid++
}

// root:0-based
//  当root设为-1时，会从0开始遍历未访问过的连通分量
func (tree *_T) Build(root int) {
	if root != -1 {
		tree.build(root, -1, 0)
		tree.markTop(root, root)
	} else {
		for i := 0; i < len(tree.Tree); i++ {
			if tree.Parent[i] == -1 {
				tree.build(i, -1, 0)
				tree.markTop(i, i)
			}
		}
	}
}

func (tree *_T) LCA(u, v int) int {
	for {
		if tree.LID[u] > tree.LID[v] {
			u, v = v, u
		}
		if tree.top[u] == tree.top[v] {
			return u
		}
		v = tree.Parent[tree.top[v]]
	}
}

// k: 0-based
//  如果不存在第k个祖先，返回-1
func (tree *_T) KthAncestor(root, k int) int {
	if k > tree.Depth[root] {
		return -1
	}
	for {
		u := tree.top[root]
		if tree.LID[root]-k >= tree.LID[u] {
			return tree.idToNode[tree.LID[root]-k]
		}
		k -= tree.LID[root] - tree.LID[u] + 1
		root = tree.Parent[u]
	}
}

// 从 from 节点跳向 to 节点,跳过 step 个节点(0-indexed)
//  返回跳到的节点,如果不存在这样的节点,返回-1
func (tree *_T) Jump(from, to, step int) int {
	if step == 1 {
		if from == to {
			return -1
		}
		if tree.IsInSubtree(to, from) {
			return tree.KthAncestor(to, tree.Depth[to]-tree.Depth[from]-1)
		}
		return tree.Parent[from]
	}
	c := tree.LCA(from, to)
	dac := tree.Depth[from] - tree.Depth[c]
	dbc := tree.Depth[to] - tree.Depth[c]
	if step > dac+dbc {
		return -1
	}
	if step <= dac {
		return tree.KthAncestor(from, step)
	}
	return tree.KthAncestor(to, dac+dbc-step)
}

func (tree *_T) CollectChild(root int) []int {
	res := []int{}
	for _, e := range tree.Tree[root] {
		next := e.to
		if next != tree.Parent[root] {
			res = append(res, next)
		}
	}
	return res
}

func (tree *_T) SubtreeSize(u int) int {
	return tree.RID[u] - tree.LID[u]
}

// child 是否在 root 的子树中 (child和root不能相等)
func (tree *_T) IsInSubtree(child, root int) bool {
	return tree.LID[root] <= tree.LID[child] && tree.LID[child] < tree.RID[root]
}

func (tree *_T) build(cur, pre, dep int) int {
	subSize, heavySize, heavySon := 1, 0, -1
	for _, e := range tree.Tree[cur] {
		next := e.to
		if next != pre {
			nextSize := tree.build(next, cur, dep+1)
			subSize += nextSize
			if nextSize > heavySize {
				heavySize, heavySon = nextSize, next
			}
		}
	}
	tree.Depth[cur] = dep
	tree.heavySon[cur] = heavySon
	tree.Parent[cur] = pre
	return subSize
}

func (tree *_T) markTop(cur, top int) {
	tree.top[cur] = top
	tree.LID[cur] = tree.timer
	tree.idToNode[tree.timer] = cur
	tree.timer++
	if tree.heavySon[cur] != -1 {
		tree.markTop(tree.heavySon[cur], top)
		for _, e := range tree.Tree[cur] {
			next := e.to
			if next != tree.heavySon[cur] && next != tree.Parent[cur] {
				tree.markTop(next, next)
			}
		}
	}
	tree.RID[cur] = tree.timer
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

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