从叶子节点开始考虑:对于每一个叶子节点,如果它本身可以被K整除,那么肯定可以直接对它和它父亲做切割。也就是说,在这种情况下,父子相连的边是一定要删掉的。
而如果不满足这一条件,这条父子相连的边最后一定会加入答案,由于需要让每一个联通块的点权之和都满足,所以如果叶子,就把叶子的点权加到父节点上即可。
自下而上的逻辑方式让DFS的实现更合理。
#include <bits/stdc++.h> using namespace std; typedef long long ll; const int N = 1000000 + 7; int to[N << 1], nxt[N << 1], head[N], wt[N << 1], tot; inline void init() { memset(head, -1, sizeof(head)), tot = 0; } inline void add(int u, int v, int val) { to[++tot] = v; nxt[tot] = head[u]; wt[tot] = val; head[u] = tot; } ll n, k, va[N], u, v, w, ans = 0; void dfs(int x, int fa) { for (int i = head[x]; ~i; i = nxt[i]) { if (to[i] == fa) continue; dfs(to[i], x); if (va[to[i]]) ans += wt[i], va[x] += va[to[i]]; } va[x] %= k; } int main() { init(); scanf("%lld%lld", &n, &k); for (int i = 1; i <= n; ++i) scanf("%lld", &va[i]), va[i] %= k; for (int i = 1; i < n; ++i) scanf("%lld%lld%lld", &u, &v, &w), add(u, v, w), add(v, u, w); dfs(1, 0); printf("%lld\n", ans); return 0; }