[SMOJ1790]蚂蚁聚会

题目描述

有 $n$ 个蚁巢，这 $n$ 个蚁巢形成一颗树形结构，第 $i$ 个蚁巢有 $a_i$ 只蚂蚁。现在蚂蚁们想举行一个大型的聚会。

但是这些蚂蚁比较懒惰，都不想走太远，每只蚂蚁最多只愿意走 $X$ 步（每一步就是走一条边）。

它们要计算：如果选择第 $i$ 个蚁巢作为举行聚会的地点，可以有多少只蚂蚁参加聚会？记该数量为 $p_i$。

你的任务就是帮助计算： $p_1, p_2, p_3, \ldots, p_n$。

输入格式 1790.in

第一行，两个整数： $n$ 和 $x$。 $1 \leq n \leq 100000$, $1 \leq x \leq 20$。

接下来有 $n-1$ 行，每行两个整数：$a$, $b$，表示蚁巢 $a$ 和蚁巢 $b$ 有一双向条边相连。

接下来有 $n$ 行，第 $i$ 行是一个整数 $a_i$。$0 \leq a_i \leq 1000$。

输出格式 1790.out

共 $n$ 行，第 $i$ 行是 $p_i$。

输入样例 1790.in

6 2
5 1
3 6
2 4
2 1
3 2
1
2
3
4
5
6

输出样例 1790.out

15
21
16
10
8
11

---

这题真的是一个好题。

题意：有一棵 $n$ 个结点的树，每个结点都有一个权值 $v_i$。现在有一个最大移动距离 $x$，且

$$v'_i=\sum v_j\text{ if }dist(i,j)\leq x$$

求 $max\left\{v'_i\right\}$。

拿样例解释一下：

样例中 $x= 2$，能到达 1 的结点有 1、2、3、4 和 5，故 $v'_1=v_1+v_2+v_3+v_4+v_5=15$。其他点同理。

做法一：暴力

枚举每一个点，做一遍 BFS，多于 $x$ 步则停止搜索，即可直接求得 $v'[]$。

时间复杂度上限：$O(n^2)$。（想象一棵树，高度为 2，除了根结点之外其他所有结点的父亲都是根结点）

得分：60。（似乎可以优化到 70，但我不会）

代码：

#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <queue>
#include <utility>
#include <vector>

using namespace std;

typedef pair <int, int> pii;

const int maxn = 1e5 + 100;

int n, x;
vector <int> g[maxn];
int a[maxn];

bool vis[maxn];
int bfs(int start) {
    memset(vis, false, sizeof vis);
    queue <pii> q;
    while (!q.empty()) q.pop();
    int ans = a[start];

    vis[start] = true;
    for (q.push(make_pair(start, 0)); !q.empty(); q.pop()) {
        pii cur = q.front();
        if (cur.second >= x) break;
        for (int i = 0; i < g[cur.first].size(); i++)
            if (!vis[g[cur.first][i]]) {
                vis[g[cur.first][i]] = true;
                q.push(make_pair(g[cur.first][i], cur.second + 1));
                ans += a[g[cur.first][i]];
            }
    }

    return ans;
}

int main(void) {
    freopen("1790.in", "r", stdin);
    freopen("1790.out", "w", stdout);

    scanf("%d%d", &n, &x);
    for (int i = 1; i < n; i++) {
        int a, b;
        scanf("%d%d", &a, &b);
        g[a].push_back(b); g[b].push_back(a);
    }
    for (int i = 1; i <= n; i++) scanf("%d", &a[i]);

    for (int i = 1; i <= n; i++) printf("%d\n", bfs(i));
    return 0;
}

做法二：树型 DP

因为移动的过程中既可以往上，也可以往下，似乎不能直接做 DP，但其实通过一点小技巧处理就可以了。

定义结点 $r$ 的深度为 $depth(r)$，我们发现，对于结点 $i$ 和 $j$（规定 $depth(i)\leq depth(j)$），它们的路径只可能有这三种情况：

当 $depth(i)<depth(j)$ 时，

1.$i$ 是 $j$ 的祖先，此时 $j$ 直接往上爬可以到达 $i$，$i$ 直接往下爬可以到达 $j$

2.$i$ 不是 $j$ 的直接祖先，但他们有公共祖先 $k$，此时 $i$ 和 $j$ 可以爬到 $k$ 之后再直接往下爬（注意 $i$ 和 $j$ 有多个公共祖先，最优路径必然取 LCA 作为中转）

当 $depth(i)=depth(j)$ 时，

3.$i$ 和 $j$ 同辈，他们有公共祖先 $k$

归纳一下，其实情况 2 和 3 可以合并。总的来看无非是：直接上下爬到达，或者通过某一个点中转。

为了方便，我们先只考虑每个结点往下的情况。不妨定义 $f[root][i]$ 为以 $root$ 为根的子树中，与 $root$ 距离在 $i$ 以内的结点权值和。

（为什么不定义为“与 $root$ 距离恰好为 $i$ 呢”？下面再讲。）

显然 $root$ 自己本身的权值一定是包含在内的，而 $f[root][i]$ 又可以分解为子问题，即所有与 $root$ 的儿子距离在 $i - 1$ 以内的结点权值和。

归纳一下，有

$$f[root][i]=weight_{root}+\sum_{j=1}^x f[son][j-1]$$

于是到目前为止我们就解决了由下往上爬的这种情况。可以确定了 $v'_i = f[i][x]$。

但是还剩下从上往下爬的情况和以某个祖先中转的情况。

这两种情况其实是一致的。

不妨这样想，当前要作为中转点的祖先 $k$ 的深度 $depth(k)$ 显然小于 $depth(i)$ 和 $depth(j)$，

于是 $i$ 到 $k$ 的距离为 $depth(i) - depth(k)$，$j$ 到 $k$ 的距离为 $depth(j) - depth(k)$，

若 $i$ 和 $j$ 可以通过 $k$ 中转而互相到达，必须满足

$$(depth(i) - depth(k)) + (depth(j) - depth(k)) \leq x$$ 即

$$depth(i) + depth(j) - 2\times depth(k) \leq x$$

由上面的推论我们可以知道，若 $i$ 和 $k$ 已经确定了，那么 $j$ 的范围也就相应确定了。

（这里就可以回答上面的疑问。对于一个确定范围的子问题，用“在 $i$ 以内”的定义方法可以在 $O(1)$ 内得到答案，如果用“恰好为 $i$”则还要再循环一遍）

而 $k$ 必然是 $i$ 的某个祖先，也就意味着我们可以从 $i$ 开始不断向上，得到每一种中转的情况。

其实如果直接用上面那个深度的式子反而还麻烦，我们不妨这样考虑吧。

有结点 $i$ 和它的某个祖先 $k$，要求两个互达的点最大距离为 $x$，由于我们是从 $i$ 不断向上找父亲的，那么其实可以知道 $depth(i) - depth(k)$。

于是也就得到了 $k$ 的另一个后裔 $j$ 与 $k$ 的距离要在 $x - (depth(i) - depth(k))$ 中。这样就可以确定子问题了，即

$$v'_i += dp[k][x - (depth(i) - depth(k))]$$

而直接从上往下爬的情况，其实相当于 $depth(i) - depth(k) = x$，所以说它们是一致的。

但是要注意，直接加会有重复。因为 $k$ 的后裔也包含了以 $i$ 为根的子树，也就可能会算到：$i$ 的后裔 $j$ 往上走到 $k$，再往下走到 $i$ 的情况，显然是不合法的。

因此要再减去重复的部分，即

$$v'_i -= f[pre(k)][x - (depth(i) - depth(k)) - 1]$$ 这里的 $pre(k)$ 为 $k$ 的儿子中 $i$ 的祖先。

得分：100。

参考代码：

#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <vector>

using namespace std;

const int maxn = 1e5 + 100;
const int maxx = 20 + 10;

int n, x;
int a[maxn];

vector <int> g[maxn];

int father[maxn]; //记录每个结点的父亲，方便后面不断往上
int dp[maxn][maxx];

void dfs(int root, int pre) { //从下往上的情况
    father[root] = pre;
    for (int i = 0; i <= x; i++) dp[root][i] = a[root]; //必定包含根结点的权值
    for (int i = 0; i < g[root].size(); i++)
        if (g[root][i] != pre) {
            dfs(g[root][i], root);
            for (int j = 1; j <= x; j++) dp[root][j] += dp[g[root][i]][j - 1];
        }
}

int main(void) {
    freopen("1790.in", "r", stdin);
    freopen("1790.out", "w", stdout);

    scanf("%d%d", &n, &x);
    for (int i = 1; i < n; i++) {
        int a, b;
        scanf("%d%d", &a, &b);
        g[a].push_back(b); g[b].push_back(a);
    }
    for (int i = 1; i <= n; i++) scanf("%d", &a[i]);

    memset(father, 0, sizeof father);
    dfs(1, 0);
    for (int i = 1; i <= n; i++) {
        int ans = dp[i][x];
        int cur = i; //当前结点
        for (int j = 1; j <= x && father[cur]; j++) { //最多往上爬 x 步，若已经爬到根结点则停止，j 即为上文所述的 depth(i) - depth(k)
            ans += dp[father[cur]][x - j]; //通过 father[cur] 中转
            if (j < x) ans -= dp[cur][x - j - 1]; //要注意对于一直往上爬 x 步的情况是不会多算的
            cur = father[cur]; //爬到它的父亲
        }
        printf("%d\n", ans);
    }
    return 0;
}