981377660LMT / ts

ts学习
6 stars 1 forks source link

ei1333 的 LinkCutTree 的模板设计与不足 #176

Open 981377660LMT opened 1 year ago

981377660LMT commented 1 year ago

https://ei1333.github.io/library/structure/lct/link-cut-tree-subtree.hpp https://ei1333.github.io/library/test/verify/yosupo-dynamic-tree-vertex-add-subtree-sum.test.cpp

981377660LMT commented 1 year ago
981377660LMT commented 1 year ago

image image

981377660LMT commented 1 year ago

完整代码:

// 动态子树和查询
// n,q=2e5
// 事件监听:
// 构造函数传入一个监听器(eventListener)在合适的时机调用监听器的回调,就可以拿到左右子树/父亲上的metadata了.

package main

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

func main() {
    // https://judge.yosupo.jp/problem/dynamic_tree_vertex_add_subtree_sum
    // 0 u v w x  删除 u-v 边, 添加 w-x 边
    // 1 p x  将 p 节点的值加上 x
    // 2 u p  对于边(u,p) 查询结点v的子树的和,p为u的父节点

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

    var n, q int
    fmt.Fscan(in, &n, &q)
    nums := make([]int, n)
    for i := 0; i < n; i++ {
        fmt.Fscan(in, &nums[i])
    }

    lct := NewLinkCutTreeSubTree(ei1333Listener{})
    vs := lct.Build(nums)

    for i := 0; i < n-1; i++ { // 连接树边
        var v1, v2 int
        fmt.Fscan(in, &v1, &v2)
        lct.LinkEdge(vs[v1], vs[v2])
    }

    for i := 0; i < q; i++ {
        var op int
        fmt.Fscan(in, &op)
        if op == 0 {
            var u, v, w, x int
            fmt.Fscan(in, &u, &v, &w, &x)
            lct.CutEdge(vs[u], vs[v])
            lct.LinkEdge(vs[w], vs[x])
        } else if op == 1 {
            var root, delta int
            fmt.Fscan(in, &root, &delta)
            lct.Set(vs[root], vs[root].key+delta)
        } else {
            var root, parent int
            fmt.Fscan(in, &root, &parent)
            lct.Evert(vs[parent])
            fmt.Fprintln(out, lct.QuerySubTree(vs[root]).csum)
        }
    }

}

type E = int
type Metadata struct {
    sum  int
    psum int
    csum int // 子树和
    lsum int
}

type ei1333Listener struct{}

func (ei1333Listener) OnToggle(cur *Metadata) {
    cur.psum, cur.csum = cur.csum, cur.psum
}

func (ei1333Listener) OnMerge(cost int, cur, parent, child *Metadata) {
    cur.sum = cost + parent.sum + child.sum + cur.lsum
    cur.psum = parent.psum + cost + cur.lsum
    cur.csum = child.csum + cost + cur.lsum
}

func (ei1333Listener) OnAdd(cur, child *Metadata) {
    cur.lsum += child.sum
}

func (ei1333Listener) OnErase(cur, child *Metadata) {
    cur.lsum -= child.sum
}

type nodeListener interface {
    OnToggle(cur *Metadata)
    OnMerge(cost int, cur, parent, child *Metadata)
    OnAdd(cur, child *Metadata)
    OnErase(cur, child *Metadata)
}

type LinkCutTreeSubTree struct {
    sm       Metadata // !ident
    listener nodeListener
}

func NewLinkCutTreeSubTree(listener nodeListener) *LinkCutTreeSubTree {
    return &LinkCutTreeSubTree{listener: listener}
}

// 各要素の値を vs[i] としたノードを生成し, その配列を返す.
func (lct *LinkCutTreeSubTree) Build(vs []E) []*treeNode {
    nodes := make([]*treeNode, len(vs))
    for i, v := range vs {
        nodes[i] = lct.Alloc(v)
    }
    return nodes
}

// 要素の値を v としたノードを生成する.
func (lct *LinkCutTreeSubTree) Alloc(key E) *treeNode {
    res := newTreeNode(key, lct.sm)
    lct.update(res)
    return res
}

// t を根に変更する.
func (lct *LinkCutTreeSubTree) Evert(t *treeNode) {
    lct.expose(t)
    lct.toggle(t)
    lct.push(t)
}

// child の親を parent にする.
//  child と parent は別の連結成分で, `child が根であること`を要求する.
func (lct *LinkCutTreeSubTree) Link(child, parent *treeNode) {
    lct.expose(child)
    lct.expose(parent)
    child.p = parent
    parent.r = child
    lct.update(parent)
}

func (lct *LinkCutTreeSubTree) LinkEdge(u, v *treeNode) {
    lct.Evert(u)
    lct.Link(u, v)
}

// child の親と child を切り離す.
func (lct *LinkCutTreeSubTree) Cut(child *treeNode) {
    lct.expose(child)
    parent := child.l
    if parent == nil {
        panic("child must not be root.")
    }
    child.l = nil
    parent.p = nil
    lct.update(child)
}

func (lct *LinkCutTreeSubTree) CutEdge(u, v *treeNode) {
    lct.Evert(u)
    lct.Cut(v)
}

// u と v の lca を返す.
//  u と v が異なる連結成分なら nullptr を返す.
//  !上記の操作は根を勝手に変えるため, 事前に Evert する必要があるかも.
func (lct *LinkCutTreeSubTree) QueryLCA(u, v *treeNode) *treeNode {
    if lct.getRoot(u) != lct.getRoot(v) {
        return nil
    }
    lct.expose(u)
    return lct.expose(v)
}

func (lct *LinkCutTreeSubTree) QueryKthAncestor(x *treeNode, k int) *treeNode {
    lct.expose(x)
    for x != nil {
        lct.push(x)
        if x.r != nil && x.r.sz > k {
            x = x.r
        } else {
            if x.r != nil {
                k -= x.r.sz
            }
            if k == 0 {
                return x
            }
            k--
            x = x.l
        }
    }
    return nil
}

func (lct *LinkCutTreeSubTree) QuerySubTree(t *treeNode) Metadata {
    lct.expose(t)
    return t.md
}

// t の値を v に変更する.
func (lct *LinkCutTreeSubTree) Set(t *treeNode, v E) *treeNode {
    lct.expose(t)
    t.key = v
    lct.update(t)
    return t
}

// t の値を返す.
func (lct *LinkCutTreeSubTree) Get(t *treeNode) E {
    return t.key
}

func (lct *LinkCutTreeSubTree) expose(t *treeNode) *treeNode {
    rp := (*treeNode)(nil)
    for cur := t; cur != nil; cur = cur.p {
        lct.splay(cur)
        if cur.r != nil {
            lct.listener.OnAdd(&cur.md, &cur.r.md)
        }
        cur.r = rp
        if cur.r != nil {
            lct.listener.OnErase(&cur.md, &cur.r.md)
        }
        lct.update(cur)
        rp = cur
    }
    lct.splay(t)
    return rp
}

func (lct *LinkCutTreeSubTree) update(t *treeNode) {
    t.sz = 1
    if t.l != nil {
        t.sz += t.l.sz
    }
    if t.r != nil {
        t.sz += t.r.sz
    }

    tmp1, tmp2 := &lct.sm, &lct.sm
    if t.l != nil {
        tmp1 = &t.l.md
    }
    if t.r != nil {
        tmp2 = &t.r.md
    }

    lct.listener.OnMerge(t.key, &t.md, tmp1, tmp2)
}

func (lct *LinkCutTreeSubTree) rotr(t *treeNode) {
    x := t.p
    y := x.p
    x.l = t.r
    if t.r != nil {
        t.r.p = x
    }
    t.r = x
    x.p = t
    lct.update(x)
    lct.update(t)
    t.p = y
    if y != nil {
        if y.l == x {
            y.l = t
        }
        if y.r == x {
            y.r = t
        }
        lct.update(y)
    }
}

func (lct *LinkCutTreeSubTree) rotl(t *treeNode) {
    x := t.p
    y := x.p
    x.r = t.l
    if t.l != nil {
        t.l.p = x
    }
    t.l = x
    x.p = t
    lct.update(x)
    lct.update(t)
    t.p = y
    if y != nil {
        if y.l == x {
            y.l = t
        }
        if y.r == x {
            y.r = t
        }
        lct.update(y)
    }
}

func (lct *LinkCutTreeSubTree) toggle(t *treeNode) {
    t.l, t.r = t.r, t.l
    lct.listener.OnToggle(&t.md)
    t.rev = !t.rev
}

func (lct *LinkCutTreeSubTree) push(t *treeNode) {
    if t.rev {
        if t.l != nil {
            lct.toggle(t.l)
        }
        if t.r != nil {
            lct.toggle(t.r)
        }
        t.rev = false
    }
}

func (lct *LinkCutTreeSubTree) splay(t *treeNode) {
    lct.push(t)
    for !t.IsRoot() {
        q := t.p
        if q.IsRoot() {
            lct.push(q)
            lct.push(t)
            if q.l == t {
                lct.rotr(t)
            } else {
                lct.rotl(t)
            }
        } else {
            r := q.p
            lct.push(r)
            lct.push(q)
            lct.push(t)
            if r.l == q {
                if q.l == t {
                    lct.rotr(q)
                    lct.rotr(t)
                } else {
                    lct.rotl(t)
                    lct.rotr(t)
                }
            } else {
                if q.r == t {
                    lct.rotl(q)
                    lct.rotl(t)
                } else {
                    lct.rotr(t)
                    lct.rotl(t)
                }
            }
        }
    }
}

func (lct *LinkCutTreeSubTree) getRoot(t *treeNode) *treeNode {
    lct.expose(t)
    for t.l != nil {
        lct.push(t)
        t = t.l
    }
    return t
}

type treeNode struct {
    key     E
    md      Metadata
    l, r, p *treeNode
    rev     bool
    sz      int
}

func newTreeNode(key E, md Metadata) *treeNode {
    return &treeNode{key: key, md: md, sz: 1}
}

func (n *treeNode) IsRoot() bool {
    return n.p == nil || (n.p.l != n && n.p.r != n)
}

func (n *treeNode) String() string {
    return fmt.Sprintf("key: %v, sum: %v, sz: %v, rev: %v", n.key, n.md, n.sz, n.rev)
}