165312025-05-06 12:14:11szilA Day in Olbiacpp17Wrong answer 32/1002.92s59976 KiB
#include <bits/stdc++.h>

using namespace std;
using ll = long long;
#define int ll

const int MAXN = 100'001;
const int K = 20;

vector<int> g[MAXN];
vector<pair<int, int>> g2[MAXN];

int tin[MAXN], tout[MAXN], up[MAXN][K], depth[MAXN], siz[MAXN], timer = 1, n;
ll ans = 0;
bool vis[MAXN];

void dfs(int u, int p = 0) {
    tin[u] = timer++;
    up[u][0] = p;
    for (int v : g[u]) {
        if (v == p) continue;
        depth[v] = depth[u] + 1;
        dfs(v, u);
    }
    tout[u] = timer++;
}

bool is_ancestor(int u, int v) {
    return tin[u] <= tin[v] && tout[v] <= tout[u];
}

void precalc() {
    for (int k = 1; k < K; k++) {
        for (int i = 1; i <= n; i++) {
            up[i][k] = up[up[i][k-1]][k-1];
        }
    }
}

int lca(int u, int v) {
    if (is_ancestor(u, v)) return u;
    if (is_ancestor(v, u)) return v;
    for (int k = K-1; k >= 0; k--) {
        if (up[u][k] && !is_ancestor(up[u][k], v)) u = up[u][k];
    }
    return up[u][0];
}

int calc_dist(int u, int v) {
    return depth[u] + depth[v] - 2 * depth[lca(u, v)];
}

void clear_virtual_tree(const vector<int> &nodes) {
    for (int i : nodes) {
        vis[i] = 0;
        g2[i].clear();
    }
}

void calc_siz(int u, int p = -1) {
    siz[u] = 1;
    for (auto [v, w] : g2[u]) {
        if (!vis[v] && v != p) {
            calc_siz(v, u);
            siz[u] += siz[v];
        }
    }
}

int find_centroid(int u, int th, int p = -1) {
    for (auto [v, w] : g2[u]) {
        if (!vis[v] && v != p && siz[v] > th) {
            return find_centroid(v, th, u);
        }
    }
    return u;
}

void dfs2(map<int, int> &cnt, int u, int p, int multi, int d, bool flush) {
    int val = multi * (u-1) - d;
    if (flush) cnt[val]++;
    else ans += cnt[-val];
    for (auto [v, w] : g2[u]) {
        if (!vis[v] && v != p) {
            dfs2(cnt, v, u, multi, d + w, flush);
        }
    }
}

void decomp(int u, int multi) {
    calc_siz(u);
    u = find_centroid(u, siz[u]/2);
    vis[u] = 1;
    map<int, int> cnt;
    cnt[(u-1)*multi]++;
    for (auto [v, w] : g2[u]) {
        if (vis[v]) continue;
        dfs2(cnt, v, u, multi, w, 0);
        dfs2(cnt, v, u, multi, w, 1);
    }
    for (auto [v, w] : g2[u]) {
        if (!vis[v]) decomp(v, multi);
    }
}

void solve() {
    cin >> n;
    for (int i = 0; i < n-1; i++) {
        int u, v; cin >> u >> v;
        u++; v++;
        g[u].emplace_back(v);
        g[v].emplace_back(u);
    }
    dfs(1);
    precalc();
    for (int c = 1; c <= n; c++) {
        vector<int> imp;
        for (int i = 1; i <= n/c; i++) {
            imp.emplace_back(i);
        }
        sort(imp.begin(), imp.end(), [&](int u, int v){
            return tin[u] < tin[v];
        });
        vector<int> nodes = imp;
        for (int i = 1; i < imp.size(); i++) {
            nodes.emplace_back(lca(imp[i-1], imp[i]));
        }
        sort(nodes.begin(), nodes.end(), [&](int u, int v){
            return tin[u] < tin[v];
        });
        nodes.erase(unique(nodes.begin(), nodes.end()), nodes.end());
        clear_virtual_tree(nodes);
        stack<int> st;
        for (int i : nodes) {
            while (!st.empty() && !is_ancestor(st.top(), i)) st.pop();
            if (!st.empty()) {
                int d = calc_dist(st.top(), i);
                g2[st.top()].emplace_back(i, d);
                g2[i].emplace_back(st.top(), d);
            }
            st.push(i);
        }
        decomp(nodes[0], c);
    }
    cout << ans << "\n";
}

signed main() {
    ios::sync_with_stdio(0); cin.tie(0);
    int t = 1; 
    // cin >> t;
    while (t--) solve();
    return 0;
}
SubtaskSumTestVerdictTimeMemory
subtask10/0
1Accepted4ms4916 KiB
2Accepted4ms4916 KiB
subtask216/16
3Accepted4ms5172 KiB
4Accepted4ms5016 KiB
5Accepted6ms5020 KiB
6Accepted6ms5200 KiB
subtask30/16
7Accepted14ms5432 KiB
8Accepted13ms5440 KiB
9Accepted14ms5612 KiB
10Accepted14ms5428 KiB
11Accepted14ms5476 KiB
12Accepted14ms5696 KiB
13Wrong answer14ms5672 KiB
14Wrong answer14ms5616 KiB
15Wrong answer13ms5636 KiB
subtask416/16
16Accepted1.825s47768 KiB
17Accepted2.375s53656 KiB
subtask50/52
18Accepted2.92s51920 KiB
19Accepted2.707s48544 KiB
20Accepted2.516s51036 KiB
21Accepted2.505s49876 KiB
22Accepted2.52s52052 KiB
23Accepted2.502s52052 KiB
24Accepted2.559s51800 KiB
25Accepted2.757s59976 KiB
26Wrong answer2.676s55960 KiB
27Accepted2.701s53656 KiB
28Accepted2.63s52888 KiB
29Accepted2.63s51916 KiB
30Accepted2.571s52100 KiB
31Wrong answer2.585s53084 KiB