CodeForces 771C Bear and Tree Jumps 树形DP

题意:

给出一棵树,一个人可以在树上跳,每次最多跳\(k(1 \leq k \leq 5)\)个点
定义\(f(s,t)\)为从顶点\(s\)跳到顶点\(t\)最少需要跳多少次
\(\sum\limits_{s<t}f(s,t)\)

分析:

注意到\(k\)很小,为了方便转移,定义:

  • \(sz(u,i)\)\(u\)的子树中\(u\)的距离模\(k\)等于\(i\)的孩子节点的个数
  • \(d(u,i)\)为从\(u\)跳到\(sz(u,i)\)对应的节点的跳数之和

转移:

  • \(sz(u,(i+1)\, mod \, k) += sz(v, i)\),其中\(v\)\(u\)的儿子节点
  • \(d(u,(i+1)\, mod \, k) += d(v, i), 0<i<k\)
  • \(d(u,1 \, mod \, k) += d(v,0) + sz(v, 0)\),因为到\(v\)距离为\(k\)的整数被的那些节点转移到\(u\)时会多跳一步

统计答案:

在递归到点\(u\)时,我们统计那些\(lca(s,t)=u\)的路径的贡献
这样的路径有两类:

  • \(u\)到其孩子节点的路径
  • 来自\(u\)的两个不同子树的孩子节点形成的路径

第一类答案很好计算,就是\(\sum d(u,i)\)

第二类答案时,假设枚举到\(u\)的某个子树\(v\),而且之前的子树的答案已经更新到到\(d(u,i),sz(u,i)\)里面
枚举\(i,j\)分别为\(v\)之前的子树距离\(u\)\(k\)\(i\)的路径,和子树\(v\)中距离\(u\)\(k\)\(j\)的路径
两边各选一段路拼起来,可以通过公式快速求得答案,细节不再赘述

最终时间复杂度为\(O(k^2n)\)

#include <cstdio>
#include <cstring>
#include <vector>
using namespace std;

typedef long long LL;

const int maxn = 200000 + 10;

vector<int> G[maxn];

int n, k, sz[maxn][5];
LL d[maxn][5], ans;

void dfs(int u, int fa) {
    int tsz[5];
    LL td[5];
    for(int v : G[u]) if(v != fa) {
        dfs(v, u);
        memset(tsz, 0, sizeof(tsz));
        for(int i = 0; i < k; i++) tsz[(i+1)%k] += sz[v][i];

        memset(td, 0, sizeof(td));
        for(int i = 1; i < k; i++) td[(i+1)%k] += d[v][i];
        td[1%k] += d[v][0] + sz[v][0];

        for(int i = 0; i < k; i++) {
            for(int j = 0; j < k; j++) {
                int delta = (i+j+k-1)/k - (i>0) - (j>0);
                ans += d[u][i] * tsz[j] + td[j] * sz[u][i];
                if(delta) ans += (LL) delta * sz[u][i] * tsz[j];
            }
        }

        for(int i = 0; i < k; i++) { d[u][i] += td[i]; sz[u][i] += tsz[i]; }
    }
    for(int i = 0; i < k; i++) ans += d[u][i];
    sz[u][0]++;
}

int main()
{
    scanf("%d%d", &n, &k);
    for(int i = 1; i < n; i++) {
        int u, v; scanf("%d%d", &u, &v);
        G[u].push_back(v);
        G[v].push_back(u);
    }

    dfs(1, 0);

    printf("%lld\n", ans);

    return 0;
}
posted @ 2017-03-19 12:46 AOQNRMGYXLMV 阅读(...) 评论(...) 编辑 收藏