D 题解 | #小美的树上染色

from functools import lru_cache
from math import sqrt

@lru_cache(None)
def dfs(u, parent, b):
    """
    param:
        b: 当前节点是否已经染色
    """
    res = 0
    t = []  # (权值，染色的贡献，不染色的贡献)
    for v in g[u]:
        if v != parent:
            t.append((ws[v], dfs(v, u, True), dfs(v, u, False)))
    sumy = sum(y for w, x, y in t)
    if b:
        # 当前已经染色，计算所有孩子节点没染色的贡献之和
        res = sumy
    else:
        # 没有染色的话，要不 不与孩子节点染色
        # 要不选取一个与之染色
        res = sumy
        wu = ws[u]
        for w, x, y in t:
            if int(sqrt(wu*w))**2 == w*wu:
                res = max(res, 2 + sumy-y+x)
    return res


n = int(input())
ws = list(map(int, input().split()))
g = [[] for _ in range(n)]
for _ in range(n-1):
    u, v = map(int, input().split())
    u, v = u-1, v-1
    g[u].append(v)
    g[v].append(u)

print(dfs(0, -1, False))