Paths on the Tree (Codeforces 1746D)
Tags: Greedy, trees

You are given a rooted tree consisting of $$$n$$$ vertices. The vertices are numbered from $$$1$$$ to $$$n$$$, and the root is the vertex $$$1$$$. You are also given a score array $$$s_1, s_2, \ldots, s_n$$$.A multiset of $$$k$$$ simple paths is called valid if the following two conditions are both true.   Each path starts from $$$1$$$.  Let $$$c_i$$$ be the number of paths covering vertex $$$i$$$. For each pair of vertices $$$(u,v)$$$ ($$$2\le u,v\le n$$$) that have the same parent, $$$|c_u-c_v|\le 1$$$ holds.  The value of the path multiset is defined as $$$\sum\limits_{i=1}^n c_i s_i$$$.It can be shown that it is always possible to find at least one valid multiset. Find the maximum value among all valid multisets.

Input: Each test contains multiple test cases. The first line contains a single integer $$$t$$$ ($$$1 \leq t \leq 10^4$$$) — the number of test cases. The description of the test cases follows.The first line of each test case contains two space-separated integers $$$n$$$ ($$$2 \le n \le 2 \cdot 10^5$$$) and $$$k$$$ ($$$1 \le k \le 10^9$$$) — the size of the tree and the required number of paths.The second line contains $$$n - 1$$$ space-separated integers $$$p_2,p_3,\ldots,p_n$$$ ($$$1\le p_i\le n$$$), where $$$p_i$$$ is the parent of the $$$i$$$-th vertex. It is guaranteed that this value describe a valid tree with root $$$1$$$.The third line contains $$$n$$$ space-separated integers $$$s_1,s_2,\ldots,s_n$$$ ($$$0 \le s_i \le 10^4$$$) — the scores of the vertices.It is guaranteed that the sum of $$$n$$$ over all test cases does not exceed $$$2 \cdot 10 ^ 5$$$.

Output: For each test case, print a single integer — the maximum value of a path multiset.

Note: In the first test case, one of optimal solutions is four paths $$$1 \to 2 \to 3 \to 5$$$, $$$1 \to 2 \to 3 \to 5$$$, $$$1 \to 4$$$, $$$1 \to 4$$$, here $$$c=[4,2,2,2,2]$$$. The value equals to $$$4\cdot 6+ 2\cdot 2+2\cdot 1+2\cdot 5+2\cdot 7=54$$$.In the second test case, one of optimal solution is three paths $$$1 \to 2 \to 3 \to 5$$$, $$$1 \to 2 \to 3 \to 5$$$, $$$1 \to 4$$$, here $$$c=[3,2,2,1,2]$$$. The value equals to $$$3\cdot 6+ 2\cdot 6+2\cdot 1+1\cdot 4+2\cdot 10=56$$$.

Sample Inputs:
2
5 4
1 2 1 3
6 2 1 5 7
5 3
1 2 1 3
6 6 1 4 10

Sample Outputs:
54
56

Solution (Python):
t = int(input())
for tidx in range(t):
    (n, k) = [int(x) for x in input().split()]
    p = [int(x) - 1 for x in input().split()]
    s = [int(x) for x in input().split()]
    children = [[] for i in range(n)]
    for i in range(n - 1):
        children[p[i]].append(i + 1)
    order = [0]
    idx = 0
    while idx < len(order):
        order.extend(children[order[idx]])
        idx += 1
    paths = [0] * n
    paths[0] = k
    for vertex in order:
        l = len(children[vertex])
        if l > 0:
            npc = paths[vertex] // l
            for c in children[vertex]:
                paths[c] = npc
    order.reverse()
    bases = [0] * n
    adds = [0] * n
    for vertex in order:
        l = len(children[vertex])
        if l == 0:
            bases[vertex] = paths[vertex] * s[vertex]
            adds[vertex] = s[vertex]
        elif l == 1:
            c = children[vertex][0]
            t1 = bases[c]
            t2 = adds[c]
            bases[vertex] = paths[vertex] * s[vertex] + t1
            adds[vertex] = s[vertex] + t2
        else:
            base = [0] * l
            add = [0] * l
            for i in range(l):
                c = children[vertex][i]
                t1 = bases[c]
                t2 = adds[c]
                base[i] = t1
                add[i] = t2
            add.sort(reverse=True)
            npaths = paths[vertex]
            bases[vertex] = sum(base) + npaths * s[vertex] + sum(add[:npaths % l])
            adds[vertex] = s[vertex] + add[npaths % l]
    ans = bases[0]
    print(ans)

Solution (Python):
import sys, threading

def DFS(v, k, boolean):
    ret = v.lable * k
    ret2 = v.lable * (k + 1)
    child_size = len(v.child)
    if child_size == 0:
        if boolean:
            return [ret, ret2]
        else:
            return [ret, 0]
    if child_size == 1:
        tmp = DFS(v.child[0], k, boolean)
        if boolean:
            return [ret + tmp[0], ret2 + tmp[1]]
        else:
            return [ret + tmp[0], 0]
    k1 = int(k / child_size)
    remind = k - k1 * child_size
    if remind == 0 and (not boolean):
        for child in v.child:
            tmp = DFS(child, k1, False)
            ret += tmp[0]
        return [ret, 0]
    tmp_list = []
    for child in v.child:
        tmp = DFS(child, k1, True)
        ret += tmp[0]
        ret2 += tmp[0]
        dif = tmp[1] - tmp[0]
        tmp_list.append(dif)
    tmp_list.sort()
    tmp_list.reverse()
    i = 0
    while i < remind:
        ret += tmp_list[i]
        ret2 += tmp_list[i]
        i += 1
    if boolean:
        ret2 += tmp_list[remind]
        return [ret, ret2]
    else:
        return [ret, 0]

class treeNode:

    def __init__(self):
        self.child = list()
        self.lable = 0

def DFS_Helper():
    t = int(input())
    for iter in range(t):
        line1 = input().split(' ')
        line2 = input().split(' ')
        line3 = input().split(' ')
        n = int(line1[0])
        k = int(line1[1])
        arr_s = [int(x) for x in line3]
        node_list = [treeNode() for i in range(0, n)]
        i = 2
        for x in line2:
            tmp = int(x)
            node_list[tmp - 1].child.append(node_list[i - 1])
            i += 1
        i = 0
        for s in arr_s:
            node_list[i].lable = s
            i += 1
        root = node_list[0]
        try:
            print(DFS(root, k, False)[0])
        except Exception as e:
            print(e)
    exit()
sys.setrecursionlimit(3 * 10 ** 5)
threading.stack_size(10 ** 8)
t1 = threading.Thread(target=DFS_Helper)
t1.start()

Solution (Python):
import sys, threading

def DFS(v, k, boolean):
    ret = v.lable * k
    ret2 = v.lable * (k + 1)
    child_size = len(v.child)
    if child_size == 0:
        if boolean:
            return [ret, ret2]
        else:
            return [ret, 0]
    if child_size == 1:
        tmp = DFS(v.child[0], k, boolean)
        if boolean:
            return [ret + tmp[0], ret2 + tmp[1]]
        else:
            return [ret + tmp[0], 0]
    k1 = int(k / child_size)
    remind = k - k1 * child_size
    if remind == 0 and (not boolean):
        for child in v.child:
            tmp = DFS(child, k1, False)
            ret += tmp[0]
        return [ret, 0]
    tmp_list = []
    for child in v.child:
        tmp = DFS(child, k1, True)
        ret += tmp[0]
        ret2 += tmp[0]
        dif = tmp[1] - tmp[0]
        tmp_list.append(dif)
    tmp_list.sort()
    tmp_list.reverse()
    i = 0
    while i < remind:
        ret += tmp_list[i]
        ret2 += tmp_list[i]
        i += 1
    if boolean:
        ret2 += tmp_list[remind]
        return [ret, ret2]
    else:
        return [ret, 0]

class treeNode:

    def __init__(self):
        self.child = list()
        self.lable = 0

def DFS_Helper():
    t = int(input())
    for iter in range(t):
        line1 = input().split(' ')
        line2 = input().split(' ')
        line3 = input().split(' ')
        n = int(line1[0])
        k = int(line1[1])
        arr_s = [int(x) for x in line3]
        node_list = [treeNode() for i in range(0, n)]
        i = 2
        for x in line2:
            tmp = int(x)
            node_list[tmp - 1].child.append(node_list[i - 1])
            i += 1
        i = 0
        for s in arr_s:
            node_list[i].lable = s
            i += 1
        root = node_list[0]
        try:
            print(DFS(root, k, False)[0])
        except Exception as e:
            print(e)
    exit()
sys.setrecursionlimit(3 * 10 ** 5)
threading.stack_size(10 ** 8)
t1 = threading.Thread(target=DFS_Helper)
t1.start()