Maximum Distributed Tree (Codeforces 1401D)
Tags: Greedy, trees

You are given a tree that consists of $$$n$$$ nodes. You should label each of its $$$n-1$$$ edges with an integer in such way that satisfies the following conditions:   each integer must be greater than $$$0$$$;  the product of all $$$n-1$$$ numbers should be equal to $$$k$$$;  the number of $$$1$$$-s among all $$$n-1$$$ integers must be minimum possible. Let's define $$$f(u,v)$$$ as the sum of the numbers on the simple path from node $$$u$$$ to node $$$v$$$. Also, let $$$\sum\limits_{i=1}^{n-1} \sum\limits_{j=i+1}^n f(i,j)$$$ be a distribution index of the tree.Find the maximum possible distribution index you can get. Since answer can be too large, print it modulo $$$10^9 + 7$$$.In this problem, since the number $$$k$$$ can be large, the result of the prime factorization of $$$k$$$ is given instead.

Input: The first line contains one integer $$$t$$$ ($$$1 \le t \le 100$$$) — the number of test cases.The first line of each test case contains a single integer $$$n$$$ ($$$2 \le n \le 10^5$$$) — the number of nodes in the tree.Each of the next $$$n-1$$$ lines describes an edge: the $$$i$$$-th line contains two integers $$$u_i$$$ and $$$v_i$$$ ($$$1 \le u_i, v_i \le n$$$; $$$u_i \ne v_i$$$) — indices of vertices connected by the $$$i$$$-th edge.Next line contains a single integer $$$m$$$ ($$$1 \le m \le 6 \cdot 10^4$$$) — the number of prime factors of $$$k$$$.Next line contains $$$m$$$ prime numbers $$$p_1, p_2, \ldots, p_m$$$ ($$$2 \le p_i < 6 \cdot 10^4$$$) such that $$$k = p_1 \cdot p_2 \cdot \ldots \cdot p_m$$$.It is guaranteed that the sum of $$$n$$$ over all test cases doesn't exceed $$$10^5$$$, the sum of $$$m$$$ over all test cases doesn't exceed $$$6 \cdot 10^4$$$, and the given edges for each test cases form a tree.

Output: Print the maximum distribution index you can get. Since answer can be too large, print it modulo $$$10^9+7$$$.

Note:  In the first test case, one of the optimal ways is on the following image:   In this case, $$$f(1,2)=1$$$, $$$f(1,3)=3$$$, $$$f(1,4)=5$$$, $$$f(2,3)=2$$$, $$$f(2,4)=4$$$, $$$f(3,4)=2$$$, so the sum of these $$$6$$$ numbers is $$$17$$$. In the second test case, one of the optimal ways is on the following image:   In this case, $$$f(1,2)=3$$$, $$$f(1,3)=1$$$, $$$f(1,4)=4$$$, $$$f(2,3)=2$$$, $$$f(2,4)=5$$$, $$$f(3,4)=3$$$, so the sum of these $$$6$$$ numbers is $$$18$$$.

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

Sample Outputs:
17
18
286

Solution (Python):
import sys
input = sys.stdin.readline

def readList():
    return list(map(int, input().split()))

def readInt():
    return int(input())

def readInts():
    return map(int, input().split())

def readStr():
    return input().strip()

def solve():
    n = readInt()
    graph = [[] for _ in range(n)]
    edges = []
    for _ in range(n - 1):
        (u, v) = readInts()
        graph[u - 1].append(v - 1)
        graph[v - 1].append(u - 1)
        edges.append((u - 1, v - 1))
    (m, p) = (readInt(), readList())
    mo = 10 ** 9 + 7
    p.sort()
    (size, st, par) = ([0] * n, [0], [-1] * n)
    par[0] = -2
    while st:
        node = st[-1]
        if size[node] == 0:
            size[node] = 1
            for nei in graph[node]:
                if par[nei] == -1:
                    par[nei] = node
                    st.append(nei)
        else:
            st.pop()
            if node > 0:
                size[par[node]] += size[node]
    t = size[0]
    arr = []
    for (u, v) in edges:
        arr.append(min(size[u], size[v]) * (t - min(size[u], size[v])))
    arr.sort()
    while len(arr) < len(p):
        arr[-1] *= p.pop()
        arr[-1] %= mo
    m = len(p)
    p = [1] * abs(n - 1 - m) + p
    ans = 0
    for i in range(n - 2, -1, -1):
        ans += arr[i] * p[i] % mo
        ans %= mo
    return ans
for _ in range(int(input())):
    print(solve())

Solution (Python):
""""""
import sys, threading

def rl():
    return [int(w) for w in sys.stdin.readline().split()]
MOD = 1000000007

def main():

    def walk(u, v):
        s = 1
        for i in adj[v]:
            if i != u:
                c = walk(v, i)
                w.append(c * (n - c))
                s += c
        return s
    for _ in range(rl()[0]):
        (n,) = rl()
        adj = [[] for _ in range(n + 1)]
        for i in range(1, n):
            (u, v) = rl()
            adj[u].append(v)
            adj[v].append(u)
        (m,) = rl()
        p = rl()
        while len(p) < n - 1:
            p.append(1)
        p.sort()
        w = []
        walk(0, 1)
        w.sort()
        r = 0
        for i in range(n - 2):
            r = (r + w[i] * p[i]) % MOD
        l = w[n - 2]
        for i in range(n - 2, max(m, n - 1)):
            l = l * p[i] % MOD
        r = (r + l) % MOD
        print(r)
sys.setrecursionlimit(10 ** 6)
threading.stack_size(10 ** 8)
t = threading.Thread(target=main)
t.start()
t.join()

Solution (Python):
import threading
import sys
(*data,) = map(int, sys.stdin.read().split()[::-1])

def inp():
    return data.pop()
modulo = 10 ** 9 + 7

def dfssize(g, v, par, size, edges):
    for ch in g[v]:
        if ch != par:
            dfssize(g, ch, v, size, edges)
            size[v] += size[ch]
            edges.append(size[ch] * (len(g) - size[ch]))
    size[v] += 1

def mult(vals):
    ans = 1
    for v in vals:
        ans = ans * v % modulo
    return ans

def solve(g, kp):
    size = [0] * len(g)
    kp.sort()
    edges = []
    dfssize(g, 0, par=-1, size=size, edges=edges)
    n = len(g)
    if len(kp) > n - 1:
        kp[n - 2:] = [mult(kp[n - 2:])]
    while len(kp) < n - 1:
        kp = [1] * (n - 1 - len(kp)) + kp
    ans = 0
    edges.sort(reverse=True)
    kp = kp[::-1]
    for (a, b) in zip(edges, kp):
        ans += a * b
        ans %= modulo
    return ans
threading.stack_size(2 ** 27)
sys.setrecursionlimit(10 ** 5 + 100)

def main():
    for _ in range(int(inp())):
        n = int(inp())
        g = [[] for _ in range(n)]
        for _ in range(n - 1):
            (src, dst) = (inp(), inp())
            src -= 1
            dst -= 1
            g[src].append(dst)
            g[dst].append(src)
        m = int(inp())
        (*kp,) = [inp() for _ in range(m)]
        print(solve(g, kp))
t = threading.Thread(target=main)
t.start()
t.join()