Heavy-Light Decompositions in Nim

abc337g.nim

# https://atcoder.jp/contests/abc337/tasks/abc337_g

include hld
import std/strutils

proc ABC337G() =
    var N = stdin.readLine.parseInt
    var tree = newGraph(N, true, N-1)
    for i in 0..<N-1:
        var uv = stdin.readLine.split.map parseInt
        tree[i] = GraphEdge(u:uv[0]-1, v:uv[1]-1)
    var hld = newHeavyLightDecomposition(tree, 0)

    var ds = newSeqWith(N+1, 0)
    var imos = newSeqWith(N+1, 0)
    proc addFwt(p : int, v : int) =
        var p = p
        p += 1
        while p <= N:
            ds[p] += v
            p += p and -p
    proc sumFwt(p : int) : int =
        var (p, s) = (p, 0)
        while p > 0:
            s += ds[p]
            p -= p and -p
        s
    for w in 0..<N:
        for c in hld.children(w):
            var (l,r) = hld.subtree(c)
            var nv = sumFwt(r) - sumFwt(l)
            imos[0] += nv; imos[l] -= nv;
            imos[r] += nv; imos[N] -= nv;
        var (l,r) = hld.subtree(w);
        var nv = sumFwt(l) + sumFwt(N) - sumFwt(r);
        imos[l] += nv; imos[r] -= nv;
        addFwt(hld.toSeq(w), 1);

    for i in 0..<N:
        imos[i+1] += imos[i]
    var ans = newSeqWith(N, 0)
    for i in 0..<N:
        ans[i] = imos[hld.toSeq(i)]
    echo ans.join " "

ABC337G()

csr_array.nim

import std/sequtils, std/algorithm
type CsrArray = ref object
    a: seq[int]
    b: seq[int]
type CsrArrayAccMiddle = ref object
    tg: CsrArray
    o: int
    p: int

proc `[]` (ca: CsrArray, i: int) : CsrArrayAccMiddle = CsrArrayAccMiddle(tg:ca, o:ca.a[i], p:ca.a[i+1])
proc `[]` (cam: CsrArrayAccMiddle, i: int) : int = cam.tg.b[cam.o + i]
proc `[]=` (cam: CsrArrayAccMiddle, i: int, v: int) = cam.tg.b[cam.o + i] = v
iterator items(cam: CsrArrayAccMiddle): int =
    for i in cam.o..<cam.p: yield cam.tg.b[i]

graph.nim

include csr_array

type GraphEdge = object
    u : int
    v : int
proc reversed(e : GraphEdge): GraphEdge = GraphEdge(u:e.v, v:e.u)

type Graph = ref object
    N: int
    undirected: bool
    edges: seq[GraphEdge]
proc newGraph(n: int, undirected: bool, m: int = 0): Graph = Graph(N:n, undirected:undirected, edges:newSeq[GraphEdge](m))

proc `[]` (g: Graph, i: int): GraphEdge = g.edges[i]
proc `[]=` (g: Graph, i: int, e: GraphEdge) = g.edges[i] = e
proc numVertices(g: Graph): int = g.N
proc numEdges(g: Graph): int = len(g.edges)
iterator items(g: Graph): GraphEdge =
    for i in g.edges: yield i

proc getAdjecencyArray(g: Graph): CsrArray =
    var n = g.numVertices()
    var m = g.numEdges()
    var c = newSeqWith(n+1, 0)
    var d = newSeq[int](m * (if g.undirected: 2 else: 1))
    for e in g:
        c[e.v] += 1
        if g.undirected: c[e.u] += 1
    for i in 0..<n: c[i+1] += c[i]
    for e in countdown(m-1, 0, 1):
        c[g[e].u] -= 1
        d[c[g[e].u]] = g[e].v
        if g.undirected:
            c[g[e].v] -= 1
            d[c[g[e].v]] = g[e].u
    CsrArray(a : c, b : d)

hld.nim

# https://nachiavivias.github.io/cp-library/cpp/tree/heavy-light-decomposition.html

include graph

type HeavyLightDecomposition = ref object
    N : int
    P : seq[int]
    PP : seq[int]
    PD : seq[int]
    D : seq[int]
    I : seq[int]
    rangeL : seq[int]
    rangeR : seq[int]

proc newHeavyLightDecomposition(tree : Graph, root : int): HeavyLightDecomposition =
    var n = tree.numVertices()
    var adj = tree.getAdjecencyArray()
    var P = newSeqWith(n, -1)
    var I = newSeqWith(n, 0)
    I[0] = root
    var iI : int = 1
    for i in 0..<n:
        var p = I[i]
        for e in adj[p]:
            if P[p] != e:
                I[iI] = e
                P[e] = p
                iI += 1
    var Z = newSeqWith(n, 1)
    var nx = newSeqWith(n, -1)
    var PP = newSeqWith(n, 0)
    for i in 0..<n:
        PP[i] = i
    for i in 1..<n:
        var p = I[n-i]
        Z[P[p]] += Z[p]
        if nx[P[p]] == -1 or Z[nx[P[p]]] < Z[p]:
            nx[P[p]] = p
    for p in I:
        if nx[p] != -1:
            PP[nx[p]] = p
    var PD = newSeqWith(n, n)
    PD[root] = 0
    var D = newSeqWith(n, 0)
    for p in I:
        if p != root:
            PP[p] = PP[PP[p]]
            PD[p] = min(PD[PP[p]], PD[P[p]]+1)
            D[p] = D[P[p]]+1
    var rangeL = newSeqWith(n, 0)
    var rangeR = newSeqWith(n, 0)
    for p in I:
        rangeR[p] = rangeL[p] + Z[p]
        var ir = rangeR[p]
        for e in adj[p]:
            if P[p] != e and e != nx[p]:
                ir -= Z[e]
                rangeL[e] = ir
        if nx[p] != -1:
            rangeL[nx[p]] = rangeL[p] + 1
    for i in 0..<n:
        I[rangeL[i]] = i
    HeavyLightDecomposition(N:n, P:P, PP:PP, D:D, I:I, rangeL:rangeL, rangeR:rangeR)

proc numVertices(hld : HeavyLightDecomposition) : int = hld.N
proc depth(hld : HeavyLightDecomposition, p : int) : int = hld.D[p]
proc toSeq(hld : HeavyLightDecomposition, vtx : int) : int = hld.rangeL[vtx]
proc toVtx(hld : HeavyLightDecomposition, seqidx : int) : int = hld.I[seqidx]
proc toSeq2In(hld : HeavyLightDecomposition, vtx : int) : int = hld.rangeL[vtx] * 2 - hld.D[vtx]
proc toSeq2Out(hld : HeavyLightDecomposition, vtx : int) : int = hld.rangeR[vtx] * 2 - hld.D[vtx] - 1
proc parentOf(hld : HeavyLightDecomposition, v : int) : int = hld.P[v]
proc heavyRootOf(hld : HeavyLightDecomposition, v : int) : int = hld.PP[v]
proc heavyChildOf(hld : HeavyLightDecomposition, v : int) : int =
    if hld.toSeq(v) == hld.N-1:
        return -1
    var cand = hld.toVtx(hld.toSeq(v) + 1)
    if hld.PP[v] == hld.PP[cand]:
        return cand
    -1
proc lca(hld : HeavyLightDecomposition, u : int, v : int) : int =
    var (u,v) = (u,v)
    if hld.PD[u] < hld.PD[v]:
        swap(u,v)
    while hld.PD[u] > hld.PD[v]:
        u = hld.P[hld.PP[u]]
    while hld.PP[u] != hld.PP[v]:
        u = hld.P[hld.PP[u]]
        v = hld.P[hld.PP[v]]
    if hld.D[u] > hld.D[v]:
        return v
    u
proc dist(hld : HeavyLightDecomposition, u : int, v : int) : int =
    hld.depth(u) + hld.depth(v) - hld.depth(hld.lca(u,v)) * 2
proc path(hld : HeavyLightDecomposition, r : int, c : int, include_root : bool, reverse_path : bool) : seq[(int,int)] =
    var (r,c) = (r,c)
    var k = hld.PD[c] - hld.PD[r] + 1
    if k <= 0:
        return @[]
    var res = newSeqWith(k, (0,0))
    for i in 0..<k-1:
        res[i] = ( hld.rangeL[hld.PP[c]], hld.rangeL[c] + 1 )
        c = hld.P[hld.PP[c]]
    if hld.PP[r] != hld.PP[c] or hld.D[r] > hld.D[c]:
        return @[]
    var root_off = if include_root: 0 else: 1
    res[^1] = ( hld.rangeL[r]+root_off, hld.rangeL[c]+1 )
    if res[^1][0] == res[^1][1]:
        discard res.pop()
        k -= 1
    if reverse_path:
        for i in 0..<k:
            res[i] = (hld.N - res[i][1], hld.N - res[i][0])
    else:
        res.reverse()
    res
proc subtree(hld : HeavyLightDecomposition, p : int) : (int, int) = ( hld.rangeL[p], hld.rangeR[p] )
proc median(hld : HeavyLightDecomposition, x : int, y : int, z : int) : int =
    hld.lca(x,y) xor hld.lca(y,z) xor hld.lca(x,z)
proc la(hld : HeavyLightDecomposition, starting : int, goal : int, d : int) : int =
    var (u, v, d) = (starting, goal, d)
    if d < 0:
        return -1
    var g = hld.lca(u,v)
    var dist0 = hld.D[u] - hld.D[g]*2 + hld.D[v]
    if dist0 < d:
        return -1
    var p = u
    if hld.D[u] - hld.D[g] < d:
        p = v
        d = dist0 - d
    while hld.D[p] - hld.D[hld.PP[p]] < d:
        d -= hld.D[p] - hld.D[hld.PP[p]] + 1
        p = hld.P[hld.PP[p]]
    hld.I[hld.rangeL[p] - d]
iterator children(hld : HeavyLightDecomposition, v : int) : int =
    var s = hld.rangeL[v] + 1
    while s < hld.rangeR[v]:
        var w = hld.toVtx(s)
        yield w
        s += hld.rangeR[w] - hld.rangeL[w]

jump_on_tree.nim

# https://judge.yosupo.jp/problem/jump_on_tree

include hld
import std/sequtils, std/strutils

proc LC_jump_on_tree() =
    var NQ = stdin.readLine.split.map parseInt
    var (N,Q) = (NQ[0], NQ[1])
    var tree = newGraph(N, true, N-1)
    for i in 0..<N-1:
        var uv = stdin.readLine.split.map parseInt
        tree[i] = GraphEdge(u:uv[0], v:uv[1])
    var hld = newHeavyLightDecomposition(tree, 0)
    for qi in 0..<Q:
        var queryData = stdin.readLine.split.map parseInt
        var (s,t,i) = (queryData[0], queryData[1], queryData[2])
        var ans = hld.la(s, t, i)
        echo ans