@keymoon氏の【全方位木DP】のGo実装

[【全方位木DP】明日使える便利な木構造のアルゴリズム](https://qiita.com/keymoon/items/2a52f1b0fb7ef67fb89e)

↑の記事における[keymoonさんによるC#のコード](https://github.com/key-moon/Library/blob/master/src/Algorithm/rerooting.csx)をGoで写経したものです。

## コード

Goにはジェネリクスがないので型 `T` を動的に変更できないのが少し気になるところですが、まぁ問題に合わせて都度書き換えるということで（セグメント木とかでもそうですが、なにかもっとスマートなプラクティスがあれば教えて下さい🙏）。

また、Stackは単なるスライシングでやっているので、もしかしたらここが速度を損ねているかもしれません。
しかしながら、これ以上コードを増やすのもどうかと思ったので、とりあえずそのままにしています。

```golang
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
}
```

## 使用例

自分が試した問題と提出を載せておきます。適宜ご参照ください。

- [EDPC V.Subtree](https://atcoder.jp/contests/dp/tasks/dp_v)
    - [提出](https://atcoder.jp/contests/dp/submissions/15728567)
- [Codeforces Round #627 (Div. 3) F.Maximum White Subtree](https://codeforces.com/contest/1324/problem/F)
    - [提出](https://codeforces.com/contest/1324/submission/89138969)
- [ABC160 F.Distributing Integers](https://atcoder.jp/contests/abc160/tasks/abc160_f)
    - [提出](https://atcoder.jp/contests/abc160/submissions/15728792)