小美的树上染色(美团2024届秋招笔试第一场编程真题)
题面
核心思想
树形DP
dp[1]表示以当前节点为根节点以及所包含的子树 且 当前节点能染色的最大染色数量
dp[0]表示以当前节点为根节点以及所包含的子树 且 当前节点不染色的最大染色数量
dp[0]好算,只用把孩子节点的dp值取max累加就好了
dp[1] 假设当前节点为cur,cur能和孩子节点nxt染色
那么dp[1] = sum(max(dp[other][0], dp[other][1]))(除了nxt的其他孩子节点的dp累加值)+ dp[nxt][0](nxt节点的dp[0]) + 2(cur和nxt染色产生的贡献)
详情看注释~
代码
import java.util.*;
public class Main {
public static void main(String[] args) {
final long MOD = (long) (1e9 + 7);
Scanner scanner = new Scanner(System.in);
int n = scanner.nextInt();
long[] value = new long[n + 1];
List<Integer>[] next = new List[n + 1];
//存放value
for(int i = 1; i <= n; i++){
value[i] = scanner.nextInt();
next[i] = new ArrayList<>();
}
//建树
for(int i = 1; i < n; i++){
int x = scanner.nextInt();
int y = scanner.nextInt();
next[x].add(y);
next[y].add(x);
}
int[] res = dpOnTheTree(1, -1, value, next);
System.out.println(Math.max(res[0], res[1]));
}
//dp[0] 表示当前节点为根 当前节点不染色的最大染色数量 dp[1]则表示当前节点染色的最大染色数量
static int[] dpOnTheTree(int cur, int pre, long[] value, List<Integer>[] next){
int[] dp = new int[2];
//存放孩子节点的dp结果
HashMap<Integer, int[]> res = new HashMap<>();
//当前节点的dp结果分两步做
//dp[0]
for(int nxt: next[cur]){
if(nxt == pre) //pre为父节点 防止走回去咯~
continue;
int[] child = dpOnTheTree(nxt, cur, value, next);
res.put(nxt, child);
// 当前节点不染色 那就是所有孩子节点的最大值和
dp[0] += Math.max(child[0], child[1]);
}
//dp[1]
for(int nxt: next[cur]){
if(nxt == pre)
continue;
long mul = value[cur] * value[nxt];
long sqrt = (long) Math.sqrt(mul);
// 可以和孩子节点染色
if(sqrt * sqrt == mul){
// dp[0] 存放的是所有孩子节点染色或不然染色的最大值和
// Math.max(res.get(nxt)[0], res.get(nxt)[1]) 取需要染色的孩子节点nxt的dp[1], dp[0]的最大值
// 以上两步计算结果 相减 就剩下了其他孩子的最大值和
// 那么当前节点和孩子节点nxt染色 dp[1] = 其他孩子节点的dp最大值的和(也就是前两步相减的结果)+ nxt这个孩子节点的dp[0] + 2
dp[1] = Math.max(dp[1], dp[0] - Math.max(res.get(nxt)[0], res.get(nxt)[1]) + res.get(nxt)[0] + 2);
}
}
return dp;
}
}
本文作者:ganyq
本文链接:https://www.cnblogs.com/ganyq/p/18111114
版权声明:本作品采用知识共享署名-非商业性使用-禁止演绎 2.5 中国大陆许可协议进行许可。
分类:
美团笔试
【推荐】国内首个AI IDE,深度理解中文开发场景,立即下载体验Trae
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步