165362025-05-06 12:23:33szilA Day in Olbiacpp17Wrong answer 0/1002.329s36012 KiB
#include <bits/stdc++.h>

using namespace std;
using ll = long long;

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+1; 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());
        nodes.erase(unique(nodes.begin(), nodes.end()), nodes.end());
        sort(nodes.begin(), nodes.end(), [&](int u, int v){
            return tin[u] < tin[v];
        });
        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";
}

int main() {
    ios::sync_with_stdio(0); cin.tie(0);
    int t = 1; 
    // cin >> t;
    while (t--) solve();
    return 0;
}
SubtaskSumTestVerdictTimeMemory
subtask10/0
1Wrong answer4ms5108 KiB
2Wrong answer4ms4916 KiB
subtask20/16
3Wrong answer4ms4960 KiB
4Wrong answer6ms5180 KiB
5Wrong answer6ms4984 KiB
6Wrong answer4ms5172 KiB
subtask30/16
7Wrong answer14ms5176 KiB
8Wrong answer10ms5172 KiB
9Wrong answer13ms5428 KiB
10Wrong answer12ms5440 KiB
11Wrong answer13ms5428 KiB
12Wrong answer13ms5620 KiB
13Wrong answer12ms5452 KiB
14Wrong answer13ms5428 KiB
15Wrong answer13ms5428 KiB
subtask40/16
16Wrong answer1.679s28528 KiB
17Runtime error74ms18860 KiB
subtask50/52
18Runtime error2.329s29404 KiB
19Runtime error93ms18344 KiB
20Runtime error101ms24492 KiB
21Runtime error82ms22184 KiB
22Runtime error79ms22836 KiB
23Wrong answer2.206s35704 KiB
24Wrong answer2.226s36012 KiB
25Runtime error81ms24492 KiB
26Runtime error81ms24420 KiB
27Runtime error76ms24492 KiB
28Runtime error75ms24492 KiB
29Runtime error75ms24492 KiB
30Runtime error75ms24296 KiB
31Runtime error71ms24492 KiB