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

很明显的树形dp问题，定义代表了当前节点选或者不选。

• 如果不选，那么：
• 如果选，在满足权值和条件的情况下，。其中d为选了其中一个孩子后的最大取值，不明白的话可以见代码。

import sys
from bisect import bisect_left
from math import inf, isqrt

# 输入加速
input = sys.stdin.readline
"""
    树形dp，当前点选或者不选
    f[i][0/1]
"""

if __name__ == '__main__':
    n = int(input())
    weight = [0] + list(map(int, input().split()))
    g = [[] for _ in range(n + 1)]
    for i in range(n - 1):
        u,v = map(int, input().split())
        g[u].append(v)
        g[v].append(u)
    f = [[0] * 2 for _ in range(n + 1)]
    def dfs(cur: int,fa: int):
        d = [-inf]
        for nxt in g[cur]:
            if nxt == fa: continue
            dfs(nxt, cur)
            w = weight[cur] * weight[nxt]
            if w == isqrt(w) ** 2:
                d.append(2 + f[nxt][0] - max(f[nxt][0], f[nxt][1]))
            f[cur][0] += max(f[nxt][0], f[nxt][1])
        f[cur][1] = f[cur][0] + max(d)
    dfs(1,0)
    print(max(f[1][0], f[1][1]))