【树上DFS】Tree and Polynomials

http://codeforces.com/gym/101372
D

push1[i][k]:所有操作1总共要让节点i下推多少系数k
push2[i][k]:所有操作2总共要让节点i上推多少系数k
sum1[i][k]:所有操作1节点i要计算多少系数k
sum2[i][k]:所有操作2节点i要计算多少系数k
遍历k从1~20跑dfs处理出所有sum1,sum2数据以及每个点的深度,最后统一计算多项式每一项
注意用memset会超时

代码:

#include 
#include 
#include 
#include 
using namespace std;
int n, k, q, e, rt;
const int maxn = (int)1e5 + 10;
typedef long long ll;
const ll mod = (ll)1e9 + 7;
struct Edge {
    int to, next;
} es[maxn];
int head[maxn];
void add(int u, int v) {
    es[e].to = v;
    es[e].next = head[u];
    head[u] = e++;
}
ll push1[maxn][30], push2[maxn][30], sum1[maxn][30], sum2[maxn][30];
int deep[maxn];
ll dfs(int u, int k, int d, ll tmp) {
    deep[u] = d;
    sum1[u][k] = (push1[u][k] + tmp) % mod;
    sum2[u][k] = push2[u][k];
    for (int i = head[u]; ~i; i = es[i].next) {
        int v = es[i].to;
        sum2[u][k] = (sum2[u][k] + dfs(v, k, d + 1, (tmp + push1[u][k]) % mod)) % mod;
    }
    return sum2[u][k];
}
ll q_pow(ll a, ll b) {
    ll ans = 1;
    while (b) {
        if (b & 1) {
            ans = ans * a % mod;
        }
        a = a * a % mod;
        b >>= 1;
    }
    return ans;
}
int main() {
    int T;
    scanf("%d", &T);
    while (T--) {
        e = 0;
        scanf("%d%d", &n, &k);
        for (int i = 0; i <= n; i++) {
            head[i] = -1;
        }
        for (int i = 1, fa; i <= n; i++) {
            scanf("%d", &fa);
            if (!fa)
                rt = i;
            else 
                add(fa, i);
            for (int j = 0; j <= k; j++) {
                push1[i][j] = push2[i][j] = 0;
            }
        }
        scanf("%d", &q);
        while (q--) {
            int op, v;
            scanf("%d%d", &op, &v);
            if (op == 1) {
                for (int i = 0; i <= k; i++) {
                    ll qs;
                    scanf("%lld", &qs);
                    push1[v][i] = (push1[v][i] + qs) % mod;
                }
            }
            else {
                for (int i = 0; i <= k; i++) {
                    ll qs;
                    scanf("%lld", &qs);
                    push2[v][i] = (push2[v][i] + qs) % mod;
                }
            }
        }
        for (int i = 0; i <= k; i++) {
            dfs(rt, i, 1, 0);
        }
        for (int i = 1; i <= n; i++) {
            ll ans = 0;
            for (int j = 0; j <= k; j++) {
                ans += ((sum1[i][j] + sum2[i][j]) % mod) * q_pow(deep[i], j) % mod;
                ans %= mod;
            }
            i == 1 ? printf("%lld", ans) : printf(" %lld", ans);
        }
        puts("");
    }
}

你可能感兴趣的:(【树上DFS】Tree and Polynomials)