GMOJ 6815. 【2020.10.06提高组模拟】树的重心

6815.【2020.10.06提高组模拟】树的重心

---

Description

---

给定一棵大小为 n 的树, 求所有大小为 k 的连通块的编号较小的重心的权值和.

Solution

---

首先, 一个点是重心的充要条件是 :

1.(siz_u geq lceil frac{k}{2} ceil).
2.对于 u 的任意儿子 v, (siz_v leq lfloor frac{k}{2} floor).

那么便可以在暴力DP上做出一点改动(暴力 设 (f_{u, i}) 表示在以 u 为根的子树中,选出包含 u 的大小为 i 的连通块的方案数),
扩展为 (f{u, i, 0/1}), 最后一维表示 u 的是否存在子树的大小超过 (lfloor frac{k}{2} floor).
转移应该十分显然.

再考虑怎么求答案, 设 (g_{u, i}) 表示以 u 为根的子树中, 选出包含 u 的大小为 i 的连通块的所有可能的重心权值和.
这时候之前对于暴力状态的扩展就体现出用处来了.
首先考虑 k 为奇数时 g 的转移 :

[ g_{u, i} = egin{cases} val_{u} imes f_{u, i - j, 0} imes f_{v, j, 0} & ({i - j leq lceil frac{k}{2} ceil, j leq lfloor frac{k}{2} floor}) \[2ex] g_{u, i - j} imes f_{v, j, 0} & ({i - j > lceil frac{k}{2} ceil}) \[2ex] g_{v, j} imes f_{u, i - j, 0} & ({j > lfloor frac{k}{2} floor + 1}) \[2ex] end{cases} ]

至于 k 为偶数, 第二条转移时需要特判 (v < u, siz_v = frac{k}{2}) 的情况, 此时应用 (val_v) 转移.

最后答案就是 (sum{g_{u, k}})

Code

#include <cstdio>
#include <cstring>
#include <algorithm>

using namespace std;

#define N 50000
#define K 500
#define Mod 1000000007

#define fo(i, x, y) for(int i = x, end_i = y; i <= end_i; i ++)
#define fd(i, x, y) for(int i = x, end_i = y; i >= end_i; i --)
#define Fo(i, u) for(int i = head[u]; i; i = edge[i].next)

void read(int &x) {
    char ch = getchar(); x = 0;
    while (ch < '0' || ch > '9') ch = getchar();
    while (ch >= '0' && ch <= '9') x = (x << 1) + (x << 3) + ch - 48, ch = getchar();
}

struct EDGE { int next, to; } edge[N << 1];

int head[N + 1], siz[N + 1], f[N + 1][K + 1][2], g[N + 1][K + 1], h[K + 1], val[N + 1], q[N + 1];

int n, m;

int cnt_edge = 1;
void Add(int u, int v) { edge[ ++ cnt_edge ] = (EDGE) { head[u], v }, head[u] = cnt_edge; }
void Link(int u, int v) { Add(u, v), Add(v, u); }

void Input() {
    int x, y;
    read(n), read(m);
    fo(i, 1, n) read(val[i]);
    fo(i, 2, n)
        read(x), read(y), Link(x, y);
}

int Min(int x, int y) { return x < y ? x : y; }
int Max(int x, int y) { return x > y ? x : y; }

void Bfs() {
    q[1] = 1; int u = 0, v = 0, flag = 0;
    for (int t = 1; t; ) {
        u = q[t], flag = 0;
        if (! siz[u]) f[u][1][0] = 1, siz[u] = 1;
        Fo(l, u) if (edge[l].to != q[t - 1]) {
            head[u] = edge[l].next;
            v = edge[l].to;
            q[ ++ t ] = v;
            flag = 1;
            break;
        }
        if (flag) continue;
        -- t;
        if (t) {
            v = q[t + 1], u = q[t];
            fo(i, 1, m) h[i] = g[u][i];
            fd(i, Min(m, siz[u] + siz[v]), 1) {
                fo(j, Max(1, i - siz[u]), Min((m >> 1), Min(siz[v], i))) {
                    if (i - j > (m >> 1) + (m & 1))
                        (h[i] += 1ll * g[u][i - j] * f[v][j][0] % Mod) %= Mod;
                    else
                        (h[i] += 1ll * ((! (m & 1) && v < u && (j << 1) == m) ? val[v] : val[u]) * f[v][j][0] % Mod * f[u][i - j][0] % Mod) %= Mod;
                    (f[u][i][0] += 1ll * f[u][i - j][0] * f[v][j][0] % Mod) %= Mod;
                    if (i - j > (m >> 1) + 1)
                        (f[u][i][1] += 1ll * f[u][i - j][1] * f[v][j][0] % Mod) %= Mod;
                }

                fo(j, Max((m >> 1) + 1, i - siz[u]), Min(siz[v], i)) {
                    (h[i] += 1ll * g[v][j] * f[u][i - j][0] % Mod) %= Mod;
                    (f[u][i][1] += 1ll * f[u][i - j][0] * (f[v][j][0] + f[v][j][1]) % Mod) %= Mod;
                }
                g[u][i] = h[i];
            }
            siz[u] += siz[v];
        }
    }
}

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

    Input();
    Bfs();

    int ans = 0;
    fo(i, 1, n) (ans += g[i][m]) %= Mod;
    printf("%d
", ans);

    return 0;
}