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

↑の記事における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
}

使用例

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