結果

問題 No.1976 Cut then Connect
ユーザー 草苺奶昔
提出日時 2023-03-16 10:31:25
言語 Go
(1.23.4)
結果
AC  
実行時間 240 ms / 2,000 ms
コード長 3,361 bytes
コンパイル時間 10,013 ms
コンパイル使用メモリ 225,256 KB
実行使用メモリ 41,160 KB
最終ジャッジ日時 2024-11-24 12:20:11
合計ジャッジ時間 13,729 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge3
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 31
権限があれば一括ダウンロードができます

ソースコード

diff # 

package main

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

func main() {

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

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

	R := NewRerootingSubtree(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, Edge{from: a, to: b, cost: 1})
	}

	// !E: (每个点处的直径,每个点到最远点的距离)
	R.ReRooting(
		func(root int) 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, edge Edge) E {
			return E{dp.dia, dp.dist + 1}
		},
	)

	res := n
	for _, e := range edges {
		u, v := e[0], e[1]
		dia1, dia2 := R.SubTree(u, v).dia, R.SubTree(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 Edge = struct{ from, to, cost int }

type Node struct {
	to, rev int
	data    Edge
}

type ReRootingSubtree struct {
	G           [][]Node
	ld, rd      [][]E
	lp, rp      []int
	e           func(root int) E
	op          func(dp1, dp2 E) E
	composition func(dp E, e Edge) E
}

func NewRerootingSubtree(n int) *ReRootingSubtree {
	res := &ReRootingSubtree{
		G:  make([][]Node, n),
		ld: make([][]E, n),
		rd: make([][]E, n),
		lp: make([]int, n),
		rp: make([]int, n),
	}
	return res
}

func (rr *ReRootingSubtree) AddEdge(u, v int, e Edge) {
	rr.G[u] = append(rr.G[u], Node{to: v, data: e})
	rr.G[v] = append(rr.G[v], Node{to: u, data: e})
}

func (rr *ReRootingSubtree) ReRooting(
	e func(root int) E,
	op func(dp1, dp2 E) E,
	composition func(dp E, e Edge) E,
) []E {
	rr.e = e
	rr.op = op
	rr.composition = composition
	n := len(rr.G)
	for i := 0; i < n; i++ {
		sort.Slice(rr.G[i], func(j, k int) bool {
			return rr.G[i][j].to < rr.G[i][k].to
		})
		rr.ld[i] = make([]E, len(rr.G[i])+1)
		rr.rd[i] = make([]E, len(rr.G[i])+1)
		for j := range rr.ld[i] {
			rr.ld[i][j] = e(i)
			rr.rd[i][j] = e(i)
		}
		rr.lp[i] = 0
		rr.rp[i] = len(rr.G[i]) - 1
	}

	for i := 0; i < n; i++ {
		for j := range rr.G[i] {
			rr.G[i][j].rev = rr.search(rr.G[rr.G[i][j].to], i)
		}
	}

	res := make([]E, n)
	for i := 0; i < n; i++ {
		res[i] = rr.dfs(i, -1)
	}
	return res
}

func (rr *ReRootingSubtree) dfs(root, eid int) E {
	for rr.lp[root] != eid && rr.lp[root] < len(rr.G[root]) {
		e := rr.G[root][rr.lp[root]]
		rr.ld[root][rr.lp[root]+1] = rr.op(rr.ld[root][rr.lp[root]], rr.composition(rr.dfs(e.to, e.rev), e.data))
		rr.lp[root]++
	}
	for rr.rp[root] != eid && rr.rp[root] >= 0 {
		e := rr.G[root][rr.rp[root]]
		rr.rd[root][rr.rp[root]] = rr.op(rr.rd[root][rr.rp[root]+1], rr.composition(rr.dfs(e.to, e.rev), e.data))
		rr.rp[root]--
	}
	if eid < 0 {
		return rr.rd[root][0]
	}
	return rr.op(rr.ld[root][eid], rr.rd[root][eid+1])
}

func (rr *ReRootingSubtree) search(vs []Node, idx int) int {
	return sort.Search(len(vs), func(i int) bool { return vs[i].to >= idx })
}

// !root 作为根节点时, 子树 v 的 dp 值
func (rr *ReRootingSubtree) SubTree(root, v int) E {
	k := rr.search(rr.G[root], v)
	return rr.composition(rr.dfs(v, rr.G[root][k].rev), rr.G[root][k].data)
}

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
}
0