Submission #1052546

---

Source Code Expand

#include <bits/stdc++.h>
#define SZ(x) ((int) (x).size())
using namespace std;

typedef long long i64;

const int NMAX = 200005;

vector<int> G[NMAX];
int values[NMAX];
int subtree[NMAX];

int h[NMAX], low[NMAX];

int total = 0;

int64_t ans = 1;

void dfs1(int node, int prev) {
    subtree[node] = values[node];
    for (int to: G[node]) {
        if (to != prev) {
            dfs1(to, node);
            h[node] = max(h[node], h[to]);
            subtree[node] += subtree[to];
        }
    }
    ++h[node];
}

void dfs2(int node, int prev, int hup) {
    int hmax1 = hup, hmax2 = 0;
    int hnode = prev;
    int hmin = NMAX;
    for (int to: G[node]) {
        if (to != prev) {
            if (h[to] > hmax1) {
                hmax2 = hmax1;
                hmax1 = h[to];
                hnode = to;
            } else if (h[to] > hmax2) {
                hmax2 = h[to];
            }
            if (subtree[to] > 0) {
                hmin = min(hmin, h[to]);
            }
        }
    }
    hmax1++; hmax2++;
    
    // cerr << node << ' ' << hmax1 << ' ' << hmax2 << endl;
    int ch;
    if (hmax1 == hmax2) {
        ch = hmax2 - 1;
    } else {
        ch = hmax2;
        if (hmax1 > hmax2 + 1) {
            ch++;
        }
    }

    if (total - subtree[node] > 0) {
        hmin = min(hmin, hup);
    }
    hmin++;

    int vis = -1;
    if (values[node] != 1) {
        if (hmin == hmax1) {
            low[node] = hmax2 + 2;
            if (hnode == prev) {
                low[node] = max(low[node], low[hnode] + 2);
            } else {
                vis = hnode;
                dfs2(hnode, node, hmax2);
                low[node] = max(low[node], low[hnode] + 2);
            }
        } else {
            low[node] = hmin;
        }
    } else {
        low[node] = 1;
    }

    ans += max(ch - low[node] + 1, 0);
    // cerr << node + 1 << ' ' << hmin << ' ' << low[node] << ' ' << ch << endl;
    
    for (int to: G[node]) {
        if (to != prev && to != vis) {
            if (h[to] == hmax1 - 1) {
                dfs2(to, node, hmax2);
            } else {
                dfs2(to, node, hmax1);
            }
        }
    }
}

int main() {
    #ifdef LOCAL_RUN
    freopen("task.in", "r", stdin);
    freopen("task.out", "w", stdout);
    //freopen("task.err", "w", stderr);
    #endif // ONLINE_JUDGE
    ios::sync_with_stdio(false);
    cin.tie(0);

    int n;
    cin >> n;

    for (int i = 1; i < n; ++i) {
        int x, y;
        cin >> x >> y;
        x--; y--;
        G[x].push_back(y);
        G[y].push_back(x);
    }

    string s;
    cin >> s;

    for (int i = 0; i < n; ++i) {
        values[i] = s[i] - '0';
        total += values[i];
    }

    dfs1(0, -1);
    dfs2(0, -1, 0);

    cout << ans << '\n';
}

Submission Info

Judge Result