↑の記事におけるkeymoonさんによるC#のコードをGoで写経したものです。
コード
Goにはジェネリクスがないので型 T を動的に変更できないのが少し気になるところですが、まぁ問題に合わせて都度書き換えるということで（セグメント木とかでもそうですが、なにかもっとスマートなプラクティスがあれば教えて下さい🙏）。
また、Stackは単なるスライシングでやっているので、もしかしたらここが速度を損ねているかもしれません。
しかしながら、これ以上コードを増やすのもどうかと思ったので、とりあえずそのままにしています。
type T int
type ReRooting struct {
NodeCount int
Identity T
Operate func(l, r T) T
OperateNode func(t T, idx int) T
Adjacents [][]int
IndexForAdjacents [][]int
Res []T
DP [][]T
}
func NewReRooting(
nodeCount int, edges [][]int, identity T, operate func(l, r T) T, operateNode func(t T, idx int) T,
) *ReRooting {
s := new(ReRooting)
s.NodeCount = nodeCount
s.Identity = identity
s.Operate = operate
s.OperateNode = operateNode
s.Adjacents = make([][]int, nodeCount)
s.IndexForAdjacents = make([][]int, nodeCount)
for _, e := range edges {
s.IndexForAdjacents[e[0]] = append(s.IndexForAdjacents[e[0]], len(s.Adjacents[e[1]]))
s.IndexForAdjacents[e[1]] = append(s.IndexForAdjacents[e[1]], len(s.Adjacents[e[0]]))
s.Adjacents[e[0]] = append(s.Adjacents[e[0]], e[1])
s.Adjacents[e[1]] = append(s.Adjacents[e[1]], e[0])
}
s.DP = make([][]T, len(s.Adjacents))
s.Res = make([]T, len(s.Adjacents))
for i := 0; i < len(s.Adjacents); i++ {
s.DP[i] = make([]T, len(s.Adjacents[i]))
}
if s.NodeCount > 1 {
s.Initialize()
} else {
s.Res[0] = s.OperateNode(s.Identity, 0)
}
return s
}
func (s *ReRooting) Query(node int) T {
return s.Res[node]
}
func (s *ReRooting) Initialize() {
parents, order := make([]int, s.NodeCount), make([]int, s.NodeCount)
// #region InitOrderedTree
index := 0
stack := []int{}
stack = append(stack, 0)
parents[0] = -1
for len(stack) > 0 {
node := stack[len(stack)-1]
stack = stack[:len(stack)-1]
order[index] = node
index++
for i := 0; i < len(s.Adjacents[node]); i++ {
adjacent := s.Adjacents[node][i]
if adjacent == parents[node] {
continue
}
stack = append(stack, adjacent)
parents[adjacent] = node
}
}
// endregion
// #region fromLeaf
for i := len(order) - 1; i >= 1; i-- {
node := order[i]
parent := parents[node]
accum := s.Identity
parentIndex := -1
for j := 0; j < len(s.Adjacents[node]); j++ {
if s.Adjacents[node][j] == parent {
parentIndex = j
continue
}
accum = s.Operate(accum, s.DP[node][j])
}
s.DP[parent][s.IndexForAdjacents[node][parentIndex]] = s.OperateNode(accum, node)
}
// endregion
// #region toLeaf
for i := 0; i < len(order); i++ {
node := order[i]
accum := s.Identity
accumsFromTail := make([]T, len(s.Adjacents[node]))
accumsFromTail[len(accumsFromTail)-1] = s.Identity
for j := len(accumsFromTail) - 1; j >= 1; j-- {
accumsFromTail[j-1] = s.Operate(s.DP[node][j], accumsFromTail[j])
}
for j := 0; j < len(accumsFromTail); j++ {
s.DP[s.Adjacents[node][j]][s.IndexForAdjacents[node][j]] = s.OperateNode(s.Operate(accum, accumsFromTail[j]), node)
accum = s.Operate(accum, s.DP[node][j])
}
s.Res[node] = s.OperateNode(accum, node)
}
// endregion
}
使用例
自分が試した問題と提出を載せておきます。適宜ご参照ください。