树中的局域重量

XMUOJ

# include <bits/stdc++.h>
using namespace std;

#define MAXSIZE 100001
vector<vector<int> > adj(MAXSIZE);
int w[MAXSIZE];

int main()
{
    int n, k, u, v;
    cin >> n >> k;
    for (int i = 0; i < n - 1; i++)
    {
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    for (int i = 1; i <= n; i++)
        cin >> w[i];
    int c[21][MAXSIZE];
    for (int i = 0; i <= n; i++)
        c[0][i] = w[i];

    for (int i = 1; i <= k; i++)
    {
        for (int j = 1; j <= n; j++)
        {
            int sum = w[j], sig;
            for (auto it : adj[j])
            {
                sig = 1;
                for (int kk = i - 1; kk >= 0; kk--)
                {
                    //cout << 't' << c[kk][it] << ' ';
                    if (sig == 1)
                        sum += c[kk][it];
                    else
                        sum -= c[kk][j]; 
                    sig = -sig;
                }
            }
            c[i][j] = sum;
        }
    }

    int flag = 0;
    for (int i = 1; i <= n; i++)
        if (!flag++)
            cout << c[k][i];
        else
            cout << endl << c[k][i];

    return 0;
}