Submission #1044235

---

Source Code Expand

#include <iostream>
#include <vector>
#include <cmath>
#include <ctime>
#include <cassert>
#include <cstdio>
#include <queue>
#include <set>
#include <map>
#include <fstream>
#include <cstdlib>
#include <string>
#include <cstring>
#include <algorithm>
#include <numeric>

#define mp make_pair
#define fi first
#define se second
#define pb push_back
#define all(x) (x).begin(), (x).end()
#define rall(x) (x).rbegin(), (x).rend()
#define forn(i, n) for (int i = 0; i < (int)(n); ++i)
#define for1(i, n) for (int i = 1; i <= (int)(n); ++i)
#define ford(i, n) for (int i = (int)(n) - 1; i >= 0; --i)
#define fore(i, a, b) for (int i = (int)(a); i <= (int)(b); ++i)

using namespace std;

typedef pair<int, int> pii;
typedef vector<int> vi;
typedef vector<vi> vvi;
typedef long long i64;
typedef vector<i64> vi64;
typedef vector<vi64> vvi64;

template<class T> bool uin(T &a, T b) { return a > b ? (a = b, true) : false; }
template<class T> bool uax(T &a, T b) { return a < b ? (a = b, true) : false; }

const int maxn = 310000;
vi e[maxn];
int d[maxn][2], dup[maxn];
int rad[maxn];
int dth[maxn];
string s;

void add(int *m, int x) {
    forn(i, 2) if (x > m[i]) swap(x, m[i]);
}

int max_exc(int *m, int x) {
    forn(i, 2) {
        if (m[i] == x) x = -1;
        else return m[i];
    }
    return -1;
}

void dfs_d(int v, int p) {
    for (int u: e[v]) {
        if (u == p) continue;
        dfs_d(u, v);
        add(d[v], max_exc(d[u], -1) + 1);
    }
}

void dfs_up(int v, int p, int D) {
    dup[v] = D;
    for (int u: e[v]) {
        if (u == p) continue;
        dfs_up(u, v, max(D, max_exc(d[v], d[u][0] + 1) + 1) + 1);
    }
}

i64 ans;

int dfs_C(int v, int p) {
    dth[v] = 0;
    int minq = s[v] == '1' ? 0 : 1e9;
    for (int u: e[v]) {
        if (u == p) continue;
        uin(minq, dfs_C(u, v) - 1);
        uax(dth[v], dth[u] + 1);
    }
    int sk = max(0, min(rad[v] - 1, dth[v] + 1) - minq + 1);
    //uin(sk, rad[v] - 1);
    ans += sk;
    minq += sk;
    if (minq == rad[v]) minq = 1e9;
    return minq;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.precision(10);
    cout << fixed;
#ifdef LOCAL_DEFINE
    freopen("input.txt", "rt", stdin);
#endif

    int n;
    cin >> n;
    forn(i, n - 1) {
        int x, y;
        cin >> x >> y;
        --x; --y;
        e[x].pb(y);
        e[y].pb(x);
    }
    cin >> s;
    forn(i, n) forn(j, 2) d[i][j] = -1;
    dfs_d(0, -1);
    dfs_up(0, -1, -1e9);
    forn(i, n) rad[i] = max(dup[i], max_exc(d[i], -1) + 1);
    int C = 0;
    forn(i, n) if (rad[i] < rad[C]) C = i;
    ans = 1;
    dfs_C(C, -1);
    cout << ans << '\n';

#ifdef LOCAL_DEFINE
    cerr << "Time elapsed: " << 1.0 * clock() / CLOCKS_PER_SEC << " s.\n";
#endif
    return 0;
}

Submission Info

Judge Result