Problem:

Range Query on a Tree II

Write a program which manipulates a weighted rooted tree $T$ with the following operations:

  
$add(v,w)$: add $w$ to all edges from the root to node $u$

  
$getSum(u)$: report the sum of weights of all edges from the root to node $u$

  The given tree $T$ consists of $n$ nodes and every node has a unique ID from $0$ to $n-1$ respectively where ID of the root is $0$.
  Note that all weights are initialized to zero.

Input

  The input is given in the following format.

$n$

$node_0$

$node_1$

$node_2$

$:$

$node_{n-1}$

$q$

$query_1$

$query_2$

$:$

$query_{q}$

  The first line of the input includes an integer $n$, the number of nodes in the tree.

  In the next $n$ lines,the information of node $i$ is given in the following format:

ki c1 c2 ... ck

  $k_i$ is the number of children of node $i$, and $c_1$ $c_2$ ... $c_{k_i}$ are node IDs of 1st, ... $k$th child of node $i$.

  In the next line, the number of queries $q$ is given. In the next $q$ lines, $i$th query is given in the following format:

0 v w

or

1 u

  The first integer represents the type of queries.'0' denotes $add(v, w)$ and '1' denotes $getSum(u)$.

Constraints

  
All the inputs are given in integers

  
$ 2 \leq n \leq 100000 $

  
$ c_j < c_{j+1} $   $( 1 \leq j \leq k-1 )$

  
$ 2 \leq q \leq 200000 $

  
$ 1 \leq u,v \leq n-1 $

  
$ 1 \leq w \leq 10000 $

Output

  For each $getSum$ query, print the sum in a line.

Sample Input 1

6
2 1 2
2 3 5
0
0
0
1 4
7
1 1
0 3 10
1 2
0 4 20
1 3
0 5 40
1 4

Sample Output 1

0
0
40
150

Sample Input 2

4
1 1
1 2
1 3
0
6
0 3 1000
0 2 1000
0 1 1000
1 1
1 2
1 3

Sample Output 2

3000
5000
6000



Sample input:
6
2 1 2
2 3 5
0
0
0
1 4
7
1 1
0 3 10
1 2
0 4 20
1 3
0 5 40
1 4


Sample output:
0
0
40
150

Submission status: Time Limit Exceeded

Solution code:
#!/usr/bin/env python3
# GRL_5_E: Tree - Range Query on Tree 2


class RAQ:
    def __init__(self, n, initial=0):
        size = 1
        while size < n:
            size *= 2
        self.size = 2*size - 1
        self.data = [initial] * self.size

    def add(self, i, j, v):
        def _add(r, lo, hi):
            if hi < i or lo > j:
                return
            elif i <= lo and hi <= j:
                self.data[r] += v
            else:
                mid = lo + (hi - lo)//2
                _add(r*2 + 1, lo, mid)
                _add(r*2 + 2, mid+1, hi)

        return _add(0, 0, self.size//2)

    def add_nested(self, ranges, v):
        def _add(r, lo, hi, n):
            while n >= 0:
                i, j = ranges[n]
                if hi < i or lo > j:
                    return
                elif (i <= lo and hi <= j):
                    self.data[r] += v
                    n -= 1
                else:
                    mid = lo + (hi - lo)//2
                    _add(r*2 + 1, lo, mid, n)
                    _add(r*2 + 2, mid+1, hi, n)
                    return

        return _add(0, 0, self.size//2, len(ranges)-1)

    def get(self, i):
        def _get(r, lo, hi, v):
            v += self.data[r]
            if lo == hi:
                return v
            mid = lo + (hi - lo)//2
            if mid >= i:
                return _get(r*2 + 1, lo, mid, v)
            else:
                return _get(r*2 + 2, mid+1, hi, v)

        return _get(0, 0, self.size//2, 0)


class Edge:
    __slots__ = ('v', 'w')

    def __init__(self, v, w):
        self.v = v
        self.w = w

    def either(self):
        return self.v

    def other(self, v):
        if v == self.v:
            return self.w
        else:
            return self.v


class Graph:
    def __init__(self, v):
        self.v = v
        self._edges = [[] for _ in range(v)]

    def add(self, e):
        self._edges[e.v].append(e)
        self._edges[e.w].append(e)

    def adj(self, v):
        return self._edges[v]

    def edges(self):
        for es in self._edges:
            for e in es:
                yield e


class PathSum2:
    """PathSum with range update.
    """
    def __init__(self, graph, root):
        self.root = root
        self.seg = RAQ(graph.v)
        self._dfs(graph, root)

    def _dfs(self, graph, root):
        visited = [False] * graph.v
        pos = [None] * graph.v
        parent = [None] * graph.v
        stack = [root]
        n = 0
        ns = []
        while stack:
            v = stack.pop()
            if not visited[v]:
                visited[v] = True
                ns.append(n)
                n += 1
                stack.append(v)
                for e in graph.adj(v):
                    w = e.other(v)
                    if not visited[w]:
                        parent[w] = v
                        stack.append(w)
            else:
                pos[v] = (ns.pop(), n - 1)

        self._pos = pos
        self._parent = parent

    def add(self, v, val):
        ps = []
        u = v
        while u != self.root:
            ps.append(self._pos[u])
            u = self._parent[u]
        self.seg.add_nested(ps, val)

    def get(self, v):
        return self.seg.get(self._pos[v][0])


def run():
    n = int(input())
    g = Graph(n)

    for i in range(n):
        k, *cs = [int(i) for i in input().split()]
        if k > 0:
            for j in cs:
                g.add(Edge(i, j))

    ps = PathSum2(g, 0)
    q = int(input())
    for _ in range(q):
        com, *args = [int(i) for i in input().split()]
        if com == 0:
            u, val = args
            ps.add(u, val)
        elif com == 1:
            u, = args
            print(ps.get(u))
        else:
            raise ValueError('invalid command')


if __name__ == '__main__':
    run()

Submission status: Accepted

Solution code:
#!/usr/bin/env python3
# GRL_5_E: Tree - Range Query on Tree 2


# Undirected Graph
class Edge:
    __slots__ = ('v', 'w')

    def __init__(self, v, w):
        self.v = v
        self.w = w

    def either(self):
        return self.v

    def other(self, v):
        if v == self.v:
            return self.w
        else:
            return self.v


class Graph:
    def __init__(self, v):
        self.v = v
        self._edges = [[] for _ in range(v)]

    def add(self, e):
        self._edges[e.v].append(e)
        self._edges[e.w].append(e)

    def adj(self, v):
        return self._edges[v]

    def edges(self):
        for es in self._edges:
            for e in es:
                yield e


class HeavyLightDecomposition:
    class _Path:
        __slots__ = ('nodes', '_index')

        def __init__(self):
            self.nodes = []
            self._index = None

        @property
        def head(self):
            return self.nodes[0]

        @property
        def size(self):
            return len(self.nodes)

        def append(self, v):
            self.nodes.append(v)

        def index(self, v):
            if self._index is None:
                self._index = {v: i for i, v in enumerate(self.nodes)}
            return self._index[v]

        def next(self, v):
            return self.nodes[self.index(v)+1]

    def __init__(self, graph, root):
        self._paths = {}
        self._parent = [root] * graph.v
        self._head = [-1] * graph.v
        self._find_paths(graph, root)

    def _find_paths(self, graph, root):
        def create_path(v):
            path = self._Path()
            u = v
            while u != -1:
                path.append(u)
                self._head[u] = v
                u = heavy[u]
            return path

        size = [1] * graph.v
        maxsize = [0] * graph.v
        heavy = [-1] * graph.v
        visited = [False] * graph.v

        stack = [root]
        while stack:
            v = stack.pop()
            if not visited[v]:
                visited[v] = True
                stack.append(v)
                for e in graph.adj(v):
                    w = e.other(v)
                    if not visited[w]:
                        self._parent[w] = v
                        stack.append(w)
            else:
                if v != root:
                    p = self._parent[v]
                    if size[v] > maxsize[p]:
                        maxsize[p] = size[v]
                        heavy[p] = v
                    size[p] += size[v]
                for e in graph.adj(v):
                    w = e.other(v)
                    if w != self._parent[v] and w != heavy[v]:
                        self._paths[w] = create_path(w)

        self._paths[root] = create_path(root)

    def parent(self, v):
        return self._parent[v]

    def paths(self):
        return self._paths.values()

    def get_path(self, v):
        return self._paths[self._head[v]]


class PathSum2:
    """PathSum with range update.
    """
    def __init__(self, graph, root):
        self.root = root
        self.hld = HeavyLightDecomposition(graph, root)
        self.rq = {p.head: RangeQuery(p.size) for p in self.hld.paths()}

    def add(self, v, val):
        u = v
        while u != self.root:
            path = self.hld.get_path(u)
            head = path.head
            if head != self.root:
                i = path.index(head)
            else:
                i = path.index(path.next(head))
            j = path.index(u)
            self.rq[head].add(i+1, j+1, val)
            u = self.hld.parent(head)

    def get(self, v):
        weight = 0
        u = v
        while u != self.root:
            path = self.hld.get_path(u)
            head = path.head
            i = path.index(head)
            j = path.index(u)
            weight += self.rq[head].sum(i+1, j+1)
            u = self.hld.parent(head)

        return weight


# Binary Indexed Tree
class BinaryIndexedTree:
    def __init__(self, n):
        self.size = n
        self.bit = [0] * (self.size+1)

    def add(self, i, v):
        while i <= self.size:
            self.bit[i] += v
            i += (i & -i)

    def sum(self, i):
        s = 0
        while i > 0:
            s += self.bit[i]
            i -= (i & -i)
        return s


class RangeQuery:
    def __init__(self, n):
        self.size = n
        self.bit1 = BinaryIndexedTree(n+1)
        self.bit2 = BinaryIndexedTree(n+1)

    def add(self, i, j, v):
        self.bit1.add(i, v * -i)
        self.bit1.add(j+1, v * (j+1))
        self.bit2.add(i, v)
        self.bit2.add(j+1, -v)

    def sum(self, i, j):
        s = self.bit1.sum(j+1) + (j+1)*self.bit2.sum(j+1)
        s -= self.bit1.sum(i) + i*self.bit2.sum(i)
        return s


def run():
    n = int(input())
    g = Graph(n)

    for i in range(n):
        k, *cs = [int(i) for i in input().split()]
        if k > 0:
            for j in cs:
                g.add(Edge(i, j))

    ps = PathSum2(g, 0)
    q = int(input())
    for _ in range(q):
        com, *args = [int(i) for i in input().split()]
        if com == 0:
            u, val = args
            ps.add(u, val)
        elif com == 1:
            u, = args
            print(ps.get(u))
        else:
            raise ValueError('invalid command')


if __name__ == '__main__':
    run()

Submission status: Accepted

Solution code:
class Tree():
    def __init__(self, n, edge):
        self.n = n
        self.tree = [[] for _ in range(n)]
        for e in edge:
            self.tree[e[0] - 1].append(e[1] - 1)
            self.tree[e[1] - 1].append(e[0] - 1)

    def setroot(self, root):
        self.root = root
        self.parent = [None for _ in range(self.n)]
        self.parent[root] = -1
        self.depth = [None for _ in range(self.n)]
        self.depth[root] = 0
        self.order = []
        self.order.append(root)
        self.size = [1 for _ in range(self.n)]
        stack = [root]
        while stack:
            node = stack.pop()
            for adj in self.tree[node]:
                if self.parent[adj] is None:
                    self.parent[adj] = node
                    self.depth[adj] = self.depth[node] + 1
                    self.order.append(adj)
                    stack.append(adj)
        for node in self.order[::-1]:
            for adj in self.tree[node]:
                if self.parent[node] == adj:
                    continue
                self.size[node] += self.size[adj]

    def heavylight_decomposition(self):
        self.order = [None for _ in range(self.n)]
        self.head = [None for _ in range(self.n)]
        self.head[self.root] = self.root
        self.next = [None for _ in range(self.n)]
        stack = [self.root]
        order = 0
        while stack:
            node = stack.pop()
            self.order[node] = order
            order += 1
            maxsize = 0
            for adj in self.tree[node]:
                if self.parent[node] == adj:
                    continue
                if maxsize < self.size[adj]:
                    maxsize = self.size[adj]
                    self.next[node] = adj
            for adj in self.tree[node]:
                if self.parent[node] == adj or self.next[node] == adj:
                    continue
                self.head[adj] = adj
                stack.append(adj)
            if self.next[node] is not None:
                self.head[self.next[node]] = self.head[node]
                stack.append(self.next[node])

    def range_hld(self, u, v, edge=False):
        res = []
        while True:
            if self.order[u] > self.order[v]: u, v = v, u
            if self.head[u] != self.head[v]:
                res.append((self.order[self.head[v]], self.order[v] + 1))
                v = self.parent[self.head[v]]
            else:
                res.append((self.order[u] + edge, self.order[v] + 1))
                return res

    def subtree_hld(self, u):
        return self.order[u], self.order[u] + self.size[u]

    def lca_hld(self, u, v):
        while True:
            if self.order[u] > self.order[v]: u, v = v, u
            if self.head[u] != self.head[v]:
                v = self.parent[self.head[v]]
            else:
                return u

class BinaryIndexedTree():  #1-indexed
    def __init__(self, n):
        self.n = n
        self.tree = [0 for _ in range(n + 1)]

    def sum(self, idx):
        res = 0
        while idx:
            res += self.tree[idx]
            idx -= idx & -idx
        return res

    def add(self, idx, x):
        while idx <= self.n:
            self.tree[idx] += x
            idx += idx & -idx

    def bisect_left(self, x):
        if x <= 0: return 0
        res, tmp = 0, 2**((self.n).bit_length() - 1)
        while tmp:
            if res + tmp <= self.n and self.tree[res + tmp] < x:
                x -= self.tree[res + tmp]
                res += tmp
            tmp >>= 1
        return res + 1

class RAQandRSQ():
    def __init__(self, n):
        self.bit1 = BinaryIndexedTree(n)
        self.bit2 = BinaryIndexedTree(n)

    def add(self, lt, rt, x):
        self.bit1.add(lt, -x * (lt - 1))
        self.bit1.add(rt, x * (rt - 1))
        self.bit2.add(lt, x)
        self.bit2.add(rt, -x)

    def sum(self, lt, rt):
        rsum = self.bit2.sum(rt - 1) * (rt - 1) + self.bit1.sum(rt - 1)
        lsum = self.bit2.sum(lt - 1) * (lt - 1) + self.bit1.sum(lt - 1)
        return rsum - lsum


import sys
input = sys.stdin.readline

from operator import add

N = int(input())
E = []

for i in range(N):
    k, *c = map(int, input().split())
    for j in range(k):
        E.append((i + 1, c[j] + 1))

t = Tree(N, E)
t.setroot(0)
t.heavylight_decomposition()

r = RAQandRSQ(N)

res = []

Q = int(input())

for _ in range(Q):
    q, *p = map(int, input().split())
    if q == 0:
        v, w = p
        for lt, rt in t.range_hld(0, v, edge=True):
            r.add(lt, rt, w)
    else:
        s = 0
        u, = p
        for lt, rt in t.range_hld(0, u, edge=True):
            s += r.sum(lt, rt)
        res.append(s)

print('\n'.join(map(str, res)))

Submission status: Accepted

Solution code:
from collections import deque
import sys
from typing import List

sys.setrecursionlimit(200000)


N = int(input())
G: List[List[int]] = [[0]] * N
for i in range(N):
    k, *c = map(lambda x: int(x), input().split())
    G[i] = c

H = [0] * N
prv = [None] * N


def dfs(v):
    s = 1
    heavy = None
    m = 0
    for w in G[v]:
        prv[w] = v
        c = dfs(w)
        if m < c:
            heavy = w
            m = c
        s += c
    H[v] = heavy
    return s


dfs(0)


SS: List[List[int]] = []
D: List[int] = []
L = [0] * N
I = [0] * N  # noqa: E741
que = deque([(0, 0)])
while que:
    v, d = que.popleft()
    S: List[int] = []
    k = len(SS)
    while v is not None:
        I[v] = len(S)
        S.append(v)
        L[v] = k
        h = H[v]
        for w in G[v]:
            if h == w:
                continue
            que.append((w, d + 1))
        v = h
    SS.append(S)
    D.append(d)


C = list(map(len, SS))
DS0 = [[0] * (c + 1) for c in C]
DS1 = [[0] * (c + 1) for c in C]


def add(K, data, k, x):
    while k <= K:
        data[k] += x
        k += k & -k


def get(K, data, k):
    s = 0
    while k:
        s += data[k]
        k -= k & -k
    return s


def query_add(v, x):
    while v is not None:
        l = L[v]  # noqa: E741
        i = I[v]
        add(C[l], DS1[l], i + 1, -x)
        add(C[l], DS1[l], 1, x)
        add(C[l], DS0[l], i + 1, x * (i + 1))
        v = prv[SS[l][0]]


def query_sum(v):
    s = - get(C[0], DS1[L[0]], 1) - get(C[0], DS0[L[0]], 1)
    while v is not None:
        l = L[v]  # noqa: E741
        i = I[v]
        s += get(C[l], DS1[l], i + 1) * (i + 1) + get(C[l], DS0[l], i + 1)
        v = prv[SS[l][0]]
    return s


Q = int(input())
ans = []
for q in range(Q):
    t, *cmd = map(lambda x: int(x), input().split())
    if t:
        ans.append(str(query_sum(cmd[0])))
    else:
        v, w = cmd
        query_add(v, w)
print("\n".join(ans))

Submission status: Accepted

Solution code:
from collections import deque

import sys
sys.setrecursionlimit(10**6)
readline = sys.stdin.readline
write = sys.stdout.write
N = int(readline())
G = [None]*N
for i in range(N):
    k, *c = map(int, readline().split())
    G[i] = c

H = [0]*N
prv = [None]*N
def dfs(v):
    s = 1; heavy = None; m = 0
    for w in G[v]:
        prv[w] = v
        c = dfs(w)
        if m < c:
            heavy = w
            m = c
        s += c
    H[v] = heavy
    return s
dfs(0)


SS = []
D = []
L = [0]*N
I = [0]*N
que = deque([(0, 0)])
while que:
    v, d = que.popleft()
    S = []
    k = len(SS)
    while v is not None:
        I[v] = len(S)
        S.append(v)
        L[v] = k
        h = H[v]
        for w in G[v]:
            if h == w:
                continue
            que.append((w, d+1))
        v = h
    SS.append(S)
    D.append(d)


C = list(map(len, SS))
DS0 = [[0]*(c+1) for c in C]
DS1 = [[0]*(c+1) for c in C]
def add(K, data, k, x):
    while k <= K:
        data[k] += x
        k += k & -k
def get(K, data, k):
    s = 0
    while k:
        s += data[k]
        k -= k & -k
    return s

def query_add(v, x):
    while v is not None:
        l = L[v]; i = I[v]
        add(C[l], DS1[l], i+1, -x)
        add(C[l], DS1[l], 1, x)
        add(C[l], DS0[l], i+1, x*(i+1))
        v = prv[SS[l][0]]

def query_sum(v):
    s = - get(C[0], DS1[L[0]], 1) - get(C[0], DS0[L[0]], 1)
    while v is not None:
        l = L[v]; i = I[v]
        s += get(C[l], DS1[l], i+1)*(i+1) + get(C[l], DS0[l], i+1)
        v = prv[SS[l][0]]
    return s


Q = int(readline())
ans = []
for q in range(Q):
    t, *cmd = map(int, readline().split())
    if t:
        ans.append(str(query_sum(cmd[0])))
    else:
        v, w = cmd
        query_add(v, w)
write("\n".join(ans))
write("\n")