【洛谷P2986】伟大的奶牛聚集

Description

给定一棵树，树上节点有点权、边有边权，求出一个点ans，使得cost最小，其中$cost=\sum\limits_{i=1}^{n}{val[i]*dis(ans, i)}$

Solution

树形dp

依旧是通过两次dfs解决

核心思想还是“二次扫描与换根法”（名词出自lyd《算法竞赛进阶指南》）

我们假定1为根并且第一遍dfs求出当ans位于1时的代价dis

显然存在$dis[i]=\sum\limits_{j\in son(i)}{dis[j]+len(i,j)*tot[j]}$,其中tot[j]表示以j为根的子树中的节点的点权之和

那么我们再进行一遍dfs实现“换根”

假设当前节点为i，并且i的父亲已经正确计算，我们记ans位于其父亲时的代价为$f[fa]$，那么我们考虑怎样通过$f[fa]$求出$f[i]$

$f[fa]$包括两部分：以i为根的子树的代价和其余部分。那么我们先将$f[fa]$减去i为根的子树的代价，然后加上i为根的子树走到i的代价，再加上剩余部分走到i的代价即可

详细实现见代码。

时间复杂度为$O(n)$

Code

 1 #include <bits/stdc++.h>
 2 using namespace std;
 3 typedef long long ll;
 4 const int maxn = 100010;
 5 inline int read() {
 6     int ret = 0, op = 1;
 7     char c = getchar();
 8     while (!isdigit(c)) {
 9         if (c == '-') op = -1; 
10         c = getchar();
11     }
12     while (isdigit(c)) {
13         ret = ret * 10 + c - '0';
14         c = getchar();
15     }
16     return ret * op;
17 }
18 struct node {
19     int nxt, to, dis;
20 } a[maxn << 1];
21 int num, head[maxn], n, val[maxn];
22 ll tot[maxn], dis[maxn], f[maxn];
23 ll sum;
24 inline void add(int from, int to, int dis) {
25     a[++num].nxt = head[from];
26     a[num].to = to;
27     a[num].dis = dis;
28     head[from] = num;
29 }
30 ll dfs(int now, int fa) {
31     for (register int i = head[now]; i; i = a[i].nxt) {
32         int to = a[i].to;
33         if (to == fa) continue ;
34         ll ret = dfs(to, now);
35         dis[now] += dis[to] + a[i].dis * ret;
36         tot[now] += ret;
37     } 
38     return tot[now] = tot[now] + val[now];
39 }
40 void dp(int now, int fa) {
41     for (register int i = head[now]; i; i = a[i].nxt) {
42         int to = a[i].to;
43         if (to == fa) continue ;
44         f[to] = (f[now] - (dis[to] + a[i].dis * tot[to])) + (a[i].dis * (sum - tot[to])) + dis[to];
45         dp(to, now);
46     }
47 }
48 int main() {
49     n = read();
50     for (register int i = 1; i <= n; ++i) val[i] = read(), sum += val[i];
51     for (register int i = 1; i < n; ++i) {
52         int x = read(), y = read(), z = read();
53         add(x, y, z); add(y, x, z);
54     }
55     dfs(1, 1);
56     f[1] = dis[1];
57     dp(1, 1);
58     ll ans = 9223372036854775806ll;
59     for (register int i = 1; i <= n; ++i) ans = min(ans, f[i]);
60     printf("%lld\n", ans);
61     return 0;
62 }

AC Code