Submission #49510742

---

Source Code Expand

// #pragma GCC optimize("Ofast,unroll-loops")
// #pragma GCC target("avx2")
// g++ "-Wl,--stack,1078749825" q1.cpp -o ab
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
typedef long double ld;
#include <ext/pb_ds/assoc_container.hpp>
using namespace __gnu_pbds;
typedef tree<ll, null_type, less<ll>, rb_tree_tag, tree_order_statistics_node_update> indexed_set;
// find_by_order(idx) and *order_of_key(val)
// for multiset use less_equal and lower_bound works as upper_bound and vice versa
#define FAST ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
#define all(x) x.begin(), x.end()
#define mxe(v) *max_element(v.begin(), v.end())
#define mne(v) *min_element(v.begin(), v.end())
#define tup(i, x) get<i>(x)
#define rem 1000000007
#define PI 3.141592653589793238462643383279502

vector<indexed_set *> subtree;
vector<ll> adj[300001];
vector<pair<ll, ll>> smaller; // smaller than self and smaller than parent in subtree

void dfs(ll i, ll parent = -1)
{
    ll largest = -1;
    vector<ll> children;
    for (ll node : adj[i])
    {
        if (node != parent)
        {
            dfs(node, i);
            children.push_back(node);
            if (largest == -1 || subtree[largest]->size() < subtree[node]->size())
            {
                largest = node;
            }
        }
    }
    if (largest == -1)
    {
        subtree[i] = new indexed_set; // new set for leaf node
    }
    else
    {
        subtree[i] = subtree[largest]; // largest sized child
    }

    for (ll child : children)
    {
        if (child == largest)
            continue;
        for (auto x : *subtree[child])
            subtree[i]->insert(x);
    }
    subtree[i]->insert(i);
    smaller[i] = {subtree[i]->order_of_key(i), subtree[i]->order_of_key(parent)};
}

vector<ll> ans;
ll sum = 0;
void reroot(ll node, vector<ll> adj[], vector<bool> &visited)
{
    visited[node] = true;
    ans[node] = sum;
    // do stuff
    for (auto x : adj[node])
    {
        if (!visited[x])
        {
            // to every child we go only through the parent while rerooting
            pair<ll, ll> p1 = smaller[node], p2 = smaller[x];
            ll osum = sum;
            sum -= p1.first;
            sum -= p2.first;
            smaller[x] = {x - 1, 0}; // correct
            ll val = x - 1 - p2.first;
            if (node < x)
            {
                val++;
            }
            smaller[node] = {node - 1 - p2.second, val};
            sum += smaller[x].first;
            sum += smaller[node].first;
            reroot(x, adj, visited);
            sum = osum;
            smaller[node] = p1;
            smaller[x] = p2;
        }
    }
}

int main()
{
    FAST;
    // freopen("input.txt", "r", stdin);
    // freopen("output.txt", "w", stdout);
    /*ll tests;
    cin>>tests;
    for (int gg=0;gg<tests;gg++)
    {}*/
    ll n;
    cin >> n;
    subtree.resize(n + 1, NULL);
    for (int i = 0; i < n - 1; i++)
    {
        ll u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    smaller.assign(n + 1, {0, 0});
    dfs(1, -1);
    for (int i = 1; i <= n; i++)
    {
        sum += smaller[i].first;
    }
    ans.assign(n + 1, 0);
    vector<bool> visit(n + 1, false);
    reroot(1, adj, visit);
    for (int i = 1; i <= n; i++)
    {
        cout << ans[i] << " ";
    }
    cout << "\n";
}

Submission Info

Judge Result