Maximizing Root (Codeforces 1778F)
Tags: Math, trees

You are given a rooted tree consisting of $$$n$$$ vertices numbered from $$$1$$$ to $$$n$$$. Vertex $$$1$$$ is the root of the tree. Each vertex has an integer value. The value of $$$i$$$-th vertex is $$$a_i$$$. You can do the following operation at most $$$k$$$ times.  Choose a vertex $$$v$$$ that has not been chosen before and an integer $$$x$$$ such that $$$x$$$ is a common divisor of the values of all vertices of the subtree of $$$v$$$. Multiply by $$$x$$$ the value of each vertex in the subtree of $$$v$$$. What is the maximum possible value of the root node $$$1$$$ after at most $$$k$$$ operations? Formally, you have to maximize the value of $$$a_1$$$.A tree is a connected undirected graph without cycles. A rooted tree is a tree with a selected vertex, which is called the root. The subtree of a node $$$u$$$ is the set of all nodes $$$y$$$ such that the simple path from $$$y$$$ to the root passes through $$$u$$$. Note that $$$u$$$ is in the subtree of $$$u$$$.

Input: The first line contains an integer $$$t$$$ ($$$1 \leq t \leq 50\,000$$$) — the number of test cases. The description of the test cases follows.The first line of each test case contains two integers $$$n$$$ and $$$k$$$ ($$$2 \leq n \leq 10^5$$$, $$$0 \leq k \leq n$$$) — the number of vertices in the tree and the number of operations.The second line contains $$$n$$$ integers $$$a_1, a_2, \ldots, a_n$$$ ($$$1 \leq a_i \leq 1000$$$), where $$$a_i$$$ denotes the value of vertex $$$i$$$.Each of the next $$$n - 1$$$ lines contains two integers $$$u_i$$$ and $$$v_i$$$ ($$$1 \leq u_i, v_i \leq n$$$, $$$u_i \neq v_i$$$), denoting the edge of the tree between vertices $$$u_i$$$ and $$$v_i$$$. It is guaranteed that the given edges form a tree.It is guaranteed that the sum of $$$n$$$ over all test cases does not exceed $$$2 \cdot 10^5$$$.

Output: For each test case, output the maximum value of the root after performing at most $$$k$$$ operations.

Note: Both examples have the same tree:  For the first test case, you can do two operations as follows:   Choose the subtree of vertex $$$4$$$ and $$$x = 2$$$.    After this operation, the node values become $$$\{24, 12, 24, 12, 12\}.$$$  Choose the subtree of vertex $$$1$$$ and $$$x = 12$$$.    After this operation, the node values become $$$\{288, 144, 288, 144, 144\}.$$$  The value of the root is $$$288$$$ and it is the maximum.For the second test case, you can do three operations as follows:   Choose the subtree of vertex $$$4$$$ and $$$x = 2$$$.    After this operation, the node values become $$$\{24, 12, 24, 12, 12\}.$$$  Choose the subtree of vertex $$$2$$$ and $$$x = 4$$$.    After this operation, the node values become $$$\{24, 48, 24, 48, 48\}.$$$  Choose the subtree of vertex $$$1$$$ and $$$x = 24$$$.    After this operation, the node values become $$$\{576, 1152, 576, 1152, 1152\}.$$$  The value of the root is $$$576$$$ and it is the maximum.

Sample Inputs:
2
5 2
24 12 24 6 12
1 2
1 3
2 4
2 5
5 3
24 12 24 6 12
1 2
1 3
2 4
2 5

Sample Outputs:
288
576

Solution (Python):
import math
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def the_sieve_of_eratosthenes(n):
    s = [1] * (n + 1)
    (s[0], s[1]) = (0, 0)
    for i in range(2, n + 1):
        if s[i]:
            for j in range(2 * i, n + 1, i):
                s[j] = 0
    return s

def divisor(i):
    s = []
    for j in range(1, int(i ** (1 / 2)) + 1):
        if i % j == 0:
            s.append(i // j)
            s.append(j)
    return sorted(set(s))

def make_graph(n, m, z):
    s = [0] * (n + 3)
    for i in range(0, 2 * m, z):
        s[x[i] + 2] += 1
    for i in range(2, n + 3):
        s[i] += s[i - 1]
    G = [0] * ((3 - z) * m)
    for i in range(0, 2 * m, z):
        j = x[i] + 1
        G[s[j]] = x[i ^ 1]
        s[j] += 1
    return (G, s)

def bfs(u):
    q = [u]
    dist = [inf] * (n + 1)
    dist[u] = 0
    for k in range(n):
        i = q[k]
        di = dist[i]
        for v in range(s[i], s[i + 1]):
            j = G[v]
            if dist[j] == inf:
                q.append(j)
                dist[j] = di + 1
    return (q, dist)

def f(u, v):
    return u * l + v
t = int(input())
ans = []
l = pow(10, 3) + 1
s = the_sieve_of_eratosthenes(l)
(d, G0, s0) = ([], [], [])
c = [-1] * l
for i in range(2, l):
    d0 = divisor(i)
    p0 = []
    for j in range(len(d0)):
        c[d0[j]] = j
        if d0[j] <= l and s[d0[j]]:
            p0.append(d0[j])
    x = []
    for j in range(len(d0)):
        u = d0[j]
        for k in p0:
            if u * k <= i and c[u * k] ^ -1:
                x.append(j)
                x.append(c[u * k])
    for j in d0:
        c[j] = -1
    (G1, s1) = make_graph(len(d0), len(x) // 2, 2)
    d.append(d0)
    G0.append(G1)
    s0.append(s1)
inf = pow(10, 9) + 1
for t0 in range(t):
    (n, k) = map(int, input().split())
    a = [inf] + list(map(int, input().split()))
    a1 = a[1]
    for i in range(2, n + 1):
        a[i] = math.gcd(a1, a[i])
    x = [0] * (2 * (n - 1))
    for i in range(n - 1):
        (u, v) = map(int, input().split())
        (x[2 * i], x[2 * i + 1]) = (u, v)
    if min(a) == 1 or not k:
        ans0 = a1
        ans.append(ans0)
        continue
    (G, s) = make_graph(n, n - 1, 1)
    (p, dist) = bfs(1)
    (d1, G1, s1) = (d[a1 - 2], G0[a1 - 2], s0[a1 - 2])
    l = len(d1)
    for i in range(l):
        c[d1[i]] = i
    dp = [inf] * ((n + 1) * l)
    for i in p[::-1]:
        ai = a[i]
        if s[i + 1] - s[i] == 1 and i ^ 1:
            for u in range(l):
                if not ai % d1[u]:
                    dp[f(i, u)] = 0
                elif not ai * ai % d1[u]:
                    dp[f(i, u)] = 1
            continue
        di = d[ai - 2]
        dp0 = [0] * l
        for v in range(s[i], s[i + 1]):
            j = G[v]
            if dist[i] < dist[j]:
                for u in range(l):
                    dp0[u] = min(dp0[u] + dp[f(j, u)], inf)
        dp1 = [inf] * l
        for u in range(l):
            g = math.gcd(d1[u], ai)
            v = c[g]
            dp1[v] = min(dp1[v], dp0[u])
            v = c[math.gcd(g * g, a1)]
            dp1[v] = min(dp1[v], dp0[u] + 1)
        for u in range(l - 1, -1, -1):
            for j in range(s1[u], s1[u + 1]):
                v = G1[j]
                dp1[u] = min(dp1[u], dp1[v])
            dp[f(i, u)] = dp1[u]
    for i in range(l):
        c[d1[i]] = -1
    ans0 = a1
    for i in range(l - 1, -1, -1):
        if dp0[i] + 1 <= k:
            ans0 = d1[i] * a1
            break
    ans.append(ans0)
sys.stdout.write('\n'.join(map(str, ans)))

Solution (Python):
import math
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def the_sieve_of_eratosthenes(n):
    s = [1] * (n + 1)
    (s[0], s[1]) = (0, 0)
    for i in range(2, n + 1):
        if s[i]:
            for j in range(2 * i, n + 1, i):
                s[j] = 0
    return s

def divisor(i):
    s = []
    for j in range(1, int(i ** (1 / 2)) + 1):
        if i % j == 0:
            s.append(i // j)
            s.append(j)
    return sorted(set(s))

def make_graph(n, m, z):
    s = [0] * (n + 3)
    for i in range(0, 2 * m, z):
        s[x[i] + 2] += 1
    for i in range(2, n + 3):
        s[i] += s[i - 1]
    G = [0] * ((3 - z) * m)
    for i in range(0, 2 * m, z):
        j = x[i] + 1
        G[s[j]] = x[i ^ 1]
        s[j] += 1
    return (G, s)

def bfs(u):
    q = [u]
    dist = [inf] * (n + 1)
    dist[u] = 0
    for k in range(n):
        i = q[k]
        di = dist[i]
        for v in range(s[i], s[i + 1]):
            j = G[v]
            if dist[j] == inf:
                q.append(j)
                dist[j] = di + 1
    return (q, dist)

def f(u, v):
    return u * l + v
t = int(input())
ans = []
l = pow(10, 3) + 1
s = the_sieve_of_eratosthenes(l)
(d, G0, s0) = ([], [], [])
c = [-1] * l
for i in range(2, l):
    d0 = divisor(i)
    p0 = []
    for j in range(len(d0)):
        c[d0[j]] = j
        if d0[j] <= l and s[d0[j]]:
            p0.append(d0[j])
    x = []
    for j in range(len(d0)):
        u = d0[j]
        for k in p0:
            if u * k <= i and c[u * k] ^ -1:
                x.append(j)
                x.append(c[u * k])
    for j in d0:
        c[j] = -1
    (G1, s1) = make_graph(len(d0), len(x) // 2, 2)
    d.append(d0)
    G0.append(G1)
    s0.append(s1)
inf = pow(10, 9) + 1
for t0 in range(t):
    (n, k) = map(int, input().split())
    a = [inf] + list(map(int, input().split()))
    a1 = a[1]
    for i in range(2, n + 1):
        a[i] = math.gcd(a1, a[i])
    x = [0] * (2 * (n - 1))
    for i in range(n - 1):
        (u, v) = map(int, input().split())
        (x[2 * i], x[2 * i + 1]) = (u, v)
    if min(a) == 1 or not k:
        ans0 = a1
        ans.append(ans0)
        continue
    (G, s) = make_graph(n, n - 1, 1)
    (p, dist) = bfs(1)
    (d1, G1, s1) = (d[a1 - 2], G0[a1 - 2], s0[a1 - 2])
    l = len(d1)
    for i in range(l):
        c[d1[i]] = i
    dp = [inf] * ((n + 1) * l)
    for i in p[::-1]:
        ai = a[i]
        if s[i + 1] - s[i] == 1 and i ^ 1:
            for u in range(l):
                if not ai % d1[u]:
                    dp[f(i, u)] = 0
                elif not ai * ai % d1[u]:
                    dp[f(i, u)] = 1
            continue
        di = d[ai - 2]
        dp0 = [0] * l
        for v in range(s[i], s[i + 1]):
            j = G[v]
            if dist[i] < dist[j]:
                for u in range(l):
                    dp0[u] = min(dp0[u] + dp[f(j, u)], inf)
        dp1 = [inf] * l
        for u in range(l):
            g = math.gcd(d1[u], ai)
            v = c[g]
            dp1[v] = min(dp1[v], dp0[u])
            v = c[math.gcd(g * g, a1)]
            dp1[v] = min(dp1[v], dp0[u] + 1)
        for u in range(l - 1, -1, -1):
            for j in range(s1[u], s1[u + 1]):
                v = G1[j]
                dp1[u] = min(dp1[u], dp1[v])
            dp[f(i, u)] = dp1[u]
    for i in range(l):
        c[d1[i]] = -1
    ans0 = a1
    for i in range(l - 1, -1, -1):
        if dp0[i] + 1 <= k:
            ans0 = d1[i] * a1
            break
    ans.append(ans0)
sys.stdout.write('\n'.join(map(str, ans)))

Solution (Python):
import math
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def the_sieve_of_eratosthenes(n):
    s = [1] * (n + 1)
    (s[0], s[1]) = (0, 0)
    for i in range(2, n + 1):
        if s[i]:
            for j in range(2 * i, n + 1, i):
                s[j] = 0
    return s

def divisor(i):
    s = []
    for j in range(1, int(i ** (1 / 2)) + 1):
        if i % j == 0:
            s.append(i // j)
            s.append(j)
    return sorted(set(s))

def make_graph(n, m, z):
    s = [0] * (n + 3)
    for i in range(0, 2 * m, z):
        s[x[i] + 2] += 1
    for i in range(2, n + 3):
        s[i] += s[i - 1]
    G = [0] * ((3 - z) * m)
    for i in range(0, 2 * m, z):
        j = x[i] + 1
        G[s[j]] = x[i ^ 1]
        s[j] += 1
    return (G, s)

def bfs(u):
    q = [u]
    dist = [inf] * (n + 1)
    dist[u] = 0
    for k in range(n):
        i = q[k]
        di = dist[i]
        for v in range(s[i], s[i + 1]):
            j = G[v]
            if dist[j] == inf:
                q.append(j)
                dist[j] = di + 1
    return (q, dist)

def f(u, v):
    return u * l + v
t = int(input())
ans = []
l = pow(10, 3) + 1
s = the_sieve_of_eratosthenes(l)
(d, G0, s0) = ([], [], [])
c = [-1] * l
for i in range(2, l):
    d0 = divisor(i)
    p0 = []
    for j in range(len(d0)):
        c[d0[j]] = j
        if d0[j] <= l and s[d0[j]]:
            p0.append(d0[j])
    x = []
    for j in range(len(d0)):
        u = d0[j]
        for k in p0:
            if u * k <= i and c[u * k] ^ -1:
                x.append(j)
                x.append(c[u * k])
    for j in d0:
        c[j] = -1
    (G1, s1) = make_graph(len(d0), len(x) // 2, 2)
    d.append(d0)
    G0.append(G1)
    s0.append(s1)
inf = pow(10, 9) + 1
for t0 in range(t):
    (n, k) = map(int, input().split())
    a = [inf] + list(map(int, input().split()))
    a1 = a[1]
    for i in range(2, n + 1):
        a[i] = math.gcd(a1, a[i])
    x = [0] * (2 * (n - 1))
    for i in range(n - 1):
        (u, v) = map(int, input().split())
        (x[2 * i], x[2 * i + 1]) = (u, v)
    if min(a) == 1 or not k:
        ans0 = a1
        ans.append(ans0)
        continue
    (G, s) = make_graph(n, n - 1, 1)
    (p, dist) = bfs(1)
    (d1, G1, s1) = (d[a1 - 2], G0[a1 - 2], s0[a1 - 2])
    l = len(d1)
    for i in range(l):
        c[d1[i]] = i
    dp = [inf] * ((n + 1) * l)
    for i in p[::-1]:
        ai = a[i]
        if s[i + 1] - s[i] == 1 and i ^ 1:
            for u in range(l):
                if not ai % d1[u]:
                    dp[f(i, u)] = 0
                elif not ai * ai % d1[u]:
                    dp[f(i, u)] = 1
            continue
        di = d[ai - 2]
        dp0 = [0] * l
        for v in range(s[i], s[i + 1]):
            j = G[v]
            if dist[i] < dist[j]:
                for u in range(l):
                    dp0[u] = min(dp0[u] + dp[f(j, u)], inf)
        dp1 = [inf] * l
        for u in range(l):
            g = math.gcd(d1[u], ai)
            v = c[g]
            dp1[v] = min(dp1[v], dp0[u])
            v = c[math.gcd(g * g, a1)]
            dp1[v] = min(dp1[v], dp0[u] + 1)
        for u in range(l - 1, -1, -1):
            for j in range(s1[u], s1[u + 1]):
                v = G1[j]
                dp1[u] = min(dp1[u], dp1[v])
            dp[f(i, u)] = dp1[u]
    for i in range(l):
        c[d1[i]] = -1
    ans0 = a1
    for i in range(l - 1, -1, -1):
        if dp0[i] + 1 <= k:
            ans0 = d1[i] * a1
            break
    ans.append(ans0)
sys.stdout.write('\n'.join(map(str, ans)))