Submission #1280783

#TimeUsernameProblemLanguageResultExecution timeMemory
1280783Bui_Quoc_CuongUnique Cities (JOI19_ho_t5)C++20
64 / 100
357 ms38696 KiB
#include <bits/stdc++.h>
using namespace std;
#define FOR(i, a, b) for (int i = a; i <= (int)b; i++)
#define FORD(i, a, b) for (int i = a; i >= (int)b; i--)
#define ll long long
#define pb push_back
#define fi first
#define se second
#define all(a) a.begin(), a.end()
#define sz(a) (int)a.size()
#define MASK(a) (1LL<<(a))
#define BIT(mask, a) ((mask >> (a))&1)
template <class A, class B> bool minimize (A &a, const B b) {if (a > b) {a = b;return true;} return false;}
template <class A, class B> bool maximize (A &a, const B b) {if (a < b) {a = b;return true;} return false;}
const int MAXN = 2e5 + 5;

int n, m;
vector <int> g[MAXN];
int a[MAXN];
int ans[MAXN];
int dp[MAXN], h[MAXN];
int root1, root2;

void init (void) {
    cin >> n >> m;
    FOR(i, 1, n - 1) {
        int u, v; cin >> u >> v;
        g[u].push_back(v); g[v].push_back(u);
    }
    FOR(i, 1, n) cin >> a[i];
}

void DFS (int u, int p) {
    for (int &v : g[u]) if (v != p) {
        h[v] = h[u] + 1;
        DFS (v, u);
    }
}

void find_dia (void) {
    DFS(1, 1);
    root1 = max_element (h + 1, h + 1 + n) - h;
    FOR(i, 1, n) h[i] = 0;
    DFS(root1, root1);
    root2 = max_element (h + 1, h + 1 + n) - h;
}

pair <int, int> st[4 * MAXN];
int lazy[4 * MAXN];

pair <int, int> merge (pair <int, int> a, pair <int, int> b) {
    pair <int, int> ans = make_pair(0, 0);
    ans.fi = min(a.fi, b.fi);
    if (ans.fi == a.fi) ans.se+= a.se;
    if (ans.fi == b.fi) ans.se+= b.se;
    return ans;
}

void apply (int id, int val) {
    st[id].first+= val;
    lazy[id]+= val;
}

void pushdown (int id) {
    if (lazy[id]) {
        apply (id << 1, lazy[id]);
        apply (id << 1 | 1, lazy[id]);
        lazy[id] = 0;
    }
}

void update (int id, int l, int r, int u, int v, int val) {
    if (l > v || r < u) return;
    if (l >= u && r <= v) {
        apply (id, val);
        return;
    }
    int mid = l + r >> 1;
    pushdown (id);
    update (id << 1, l, mid, u, v, val);
    update (id << 1 | 1, mid + 1, r, u, v, val);
    st[id] = merge (st[id << 1], st[id << 1 | 1]);
}

pair <int, int> get (int id, int l, int r, int u, int v) {
    if (l > v || r < u) return make_pair (2e9, 0);
    if (l >= u && r <= v) return st[id];
    pushdown (id);
    int mid = l + r >> 1;
    return merge(get (id << 1, l, mid, u, v), get (id << 1 | 1, mid + 1, r, u, v));
}

void build (int id, int l, int r) {
    if (l == r) {
        st[id].first = 0;
        st[id].second = 1;
        return;
    }
    int mid = (l + r) >> 1;
    build (id << 1, l, mid);
    build (id << 1 | 1, mid + 1, r);
    st[id] = merge (st[id << 1], st[id << 1 | 1]);
}

void pre_dp (int u, int p) {
    for (int &v : g[u]) if (v != p) {
        h[v] = h[u] + 1;
        pre_dp (v, u);
        dp[u] = max(dp[u], dp[v] + 1);
    }
}

void dfs (int u, int p) {
    if (h[u] > dp[u]) {
        pair <int, int> res = get (1, 1, n, 1, h[u] - dp[u]);
        if (res.first == 0) {
            maximize (ans[u], res.second);
        }
    }
    pair <int, int> best0 = make_pair(- 1, 0);
    pair <int, int> best1 = make_pair(- 1, 0);

    for (int &v : g[u]) if (v != p) {
        if (best0.fi < dp[v]) {
            best1 = best0;
            best0 = make_pair(dp[v], v);
        } else if (best1.fi < dp[v]) {
            best1 = make_pair(dp[v], v);
        }
    }

    for (int &v : g[u]) if (v != p) {
        if (v == best0.second) {
            update (1, 1, n, h[u] - best1.first, h[u], 1);
            dfs (v, u);
            update (1, 1, n, h[u] - best1.first, h[u], - 1);
        } else {
            update (1, 1, n, h[u] - best0.first, h[u], 1);
            dfs (v, u);
            update (1, 1, n, h[u] - best0.first, h[u], - 1);
        }
    }
}

void process (void) {
    find_dia();
    FOR(i, 1, n) h[i] = dp[i] = 0;
    build (1, 1, n);
    pre_dp (root1, root1);
    dfs (root1, root1);

    FOR(i, 1, n) h[i] = dp[i] = 0;
    build (1, 1, n);
    FOR(i, 1, 4 * n) lazy[i] = 0;
    pre_dp (root2, root2);
    dfs (root2, root2);

    FOR(i, 1, n) {
        if (m == 1) {
            cout << (ans[i] > 0 ? 1 : 0) << "\n";
        } else cout << ans[i] << "\n";
    }
}

int main (void) {
    ios_base::sync_with_stdio(0); cin.tie(0);
    #define kieuoanh "kieuoanh"
    if (fopen(kieuoanh".inp", "r")) {
        freopen(kieuoanh".inp", "r", stdin); freopen(kieuoanh".out", "w", stdout);
    }
    init();
    process();
    return 0;
}

Compilation message (stderr)

joi2019_ho_t5.cpp: In function 'int main()':
joi2019_ho_t5.cpp:169:16: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  169 |         freopen(kieuoanh".inp", "r", stdin); freopen(kieuoanh".out", "w", stdout);
      |         ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~
joi2019_ho_t5.cpp:169:53: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  169 |         freopen(kieuoanh".inp", "r", stdin); freopen(kieuoanh".out", "w", stdout);
      |                                              ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...