結果

問題 No.3194 Do Optimize Your Solution
コンテスト
ユーザー noya2
提出日時 2025-06-25 16:28:44
言語 C++23
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 950 ms / 3,000 ms
コード長 6,009 bytes
コンパイル時間 4,445 ms
コンパイル使用メモリ 315,976 KB
実行使用メモリ 84,132 KB
最終ジャッジ日時 2025-06-27 20:54:52
合計ジャッジ時間 14,099 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 17
権限があれば一括ダウンロードができます

ソースコード

diff # 

#pragma GCC optimize("O3")
#include<bits/stdc++.h>
#define rep(i,s,n) for (int i = (int)(s); i < (int)(n); i++)
using namespace std;
using ull = unsigned long long;
const int mx = 18;

int input(){
    int x = 0;
    char c;
    while ((c = getchar_unlocked() ^ 48) < 10){
        x = x * 10 + c;
    }
    return x;
}

template<class E>
struct csr {
    csr (int _n, int m) : n(_n), back(0){
        start.resize(m);
        elist.resize(m);
    }
    void reset(){
        back = 0;
    }
    void add(int idx, E elem){
        start[back] = idx;
        elist[back] = elem;
        back++;
    }
    void build(){
        int m = start.size();
        std::vector<E> nelist(m);
        std::vector<int> nstart(n + 2, 0);
        for (int i = 0; i < m; i++){
            nstart[start[i] + 2]++;
        }
        for (int i = 1; i < n; i++){
            nstart[i + 2] += nstart[i + 1];
        }
        for (int i = 0; i < m; i++){
            nelist[nstart[start[i] + 1]++] = elist[i];
        }
        swap(elist,nelist);
        for (int i = 0; i < n + 1; i++){
            start[i] = nstart[i];
        }
    }
    const auto operator[](int idx) const {
        return std::ranges::subrange(elist.begin()+start[idx],elist.begin()+start[idx+1]);
    }
    auto operator[](int idx){
        return std::ranges::subrange(elist.begin()+start[idx],elist.begin()+start[idx+1]);
    }
    size_t size() const {
        return n;
    }
    int n, back;
    std::vector<int> start;
    std::vector<E> elist;
};

int main(){
    cin.tie(0)->sync_with_stdio(0);
    int n = input();
    csr<int> a(n,2*(n-1));
    rep(i,0,n-1){
        int u = input(); u--;
        int v = input(); v--;
        a.add(u,v);
        a.add(v,u);
    }
    a.build();

    vector<bool> done(n,false);
    vector<int> par(n);
    vector<array<int,mx>> dist(n);
    struct dat {
        int cnt, ecnt;
        ull sum, esum;
    };
    vector<dat> dp(n);
    vector<int16_t> dep(n);
    auto cd = [&](auto _cd, int one, int sz, int16_t recdep) -> pair<int,ull> {
        int ctr = -1, ctrpar = -1;
        ull dsum = 0;
        auto subsz = [&](auto _subsz, int v, int f, int d) -> void {
            dp[v].cnt = 1;
            dist[v][recdep] = d;
            dsum += d;
            for (int u : a[v]){
                if (done[u]) continue;
                if (u == f) continue;
                _subsz(_subsz,u,v,d+1);
                dp[v].cnt += dp[u].cnt;
            }
            if (dp[v].cnt*2 >= sz){
                if (ctr == -1){
                    ctr = v;
                    ctrpar = f;
                }
            }
        };
        subsz(subsz,one,-1,1);
        dep[ctr] = recdep;
        done[ctr] = true;
        ull ssum = 0;
        for (int none : a[ctr]){
            if (done[none]) continue;
            int c0 = (none == ctrpar ? sz - dp[ctr].cnt : dp[none].cnt);
            auto [ch, s0] = _cd(_cd,none,c0,recdep+1);
            par[ch] = ctr;
            dp[ch].ecnt = c0;
            dp[ch].esum = s0;
            ssum += s0;
        }
        dp[ctr].cnt = sz;
        dp[ctr].sum = ssum;
        return {ctr, dsum};
    };
    cd(cd,0,n,0);

    // sum01 = sum[u : 0][v : 1] dist(u,v)
    //       = sum[ctr][u : 0][v : 1] (dist_from_ctr[u] + dist_from_ctr[v])
    //       = sum[ctr] ( cnt1[ctr] * sum0[ctr] + cnt0[ctr] * sum1[ctr] )
    ull sum01 = 0;
    auto turn_on = [&](int v){
        sum01 += dp[v].sum;
        dp[v].cnt -= 2;
        int16_t recdep = dep[v];
        int c = v;
        auto distv = dist[v];
        while (recdep > 0){
            int p = par[c];
            recdep--;
            int d = distv[recdep+1];
            sum01 += dp[p].sum - dp[c].esum;
            sum01 += ull(dp[p].cnt - dp[c].ecnt) * d;
            dp[c].ecnt -= 2;
            dp[c].esum -= 2*d;
            dp[p].cnt -= 2;
            dp[p].sum -= 2*d;
            c = p;
        }
    };
    auto turn_off = [&](int v){
        sum01 -= dp[v].sum;
        dp[v].cnt += 2;
        int16_t recdep = dep[v];
        int c = v;
        auto distv = dist[v];
        while (recdep > 0){
            int p = par[c];
            recdep--;
            int d = distv[recdep+1];
            sum01 -= dp[p].sum - dp[c].esum;
            sum01 -= ull(dp[p].cnt - dp[c].ecnt) * d;
            dp[c].ecnt += 2;
            dp[c].esum += 2*d;
            dp[p].cnt += 2;
            dp[p].sum += 2*d;
            c = p;
        }
    };
    a.reset();
    rep(i,0,n-1){
        int u = input(); u--;
        int v = input(); v--;
        a.add(u,v);
        a.add(v,u);
    }
    a.build();
    ull ans = 0;
    vector<int> sub(n), down(n), tour(n), depsum(n);
    int t = 0;
    auto subsz = [&](auto _subsz, int v, int f) -> void {
        sub[v] = 1;
        depsum[v] = dep[v] + 1;
        tour[t] = v;
        down[v] = t++;
        for (int u : a[v]){
            if (u == f) continue;
            _subsz(_subsz,u,v);
            sub[v] += sub[u];
            depsum[v] += depsum[u];
        }
    };
    subsz(subsz,0,-1);
    auto dfs = [&](auto _dfs, int v, int f, bool top) -> void {
        int heavy = -1, ch = -1;
        for (int u : a[v]){
            if (u == f) continue;
            if (heavy < depsum[u]){
                heavy = depsum[u];
                ch = u;
            }
        }
        for (int u : a[v]){
            if (u == f) continue;
            if (u == ch) continue;
            _dfs(_dfs,u,v,true);
        }
        if (ch != -1){
            _dfs(_dfs,ch,v,false);
        }
        if (v == 0) return ;
        for (int u : a[v]){
            if (u == f) continue;
            if (u == ch) continue;
            for (int i = down[u]; i < down[u]+sub[u]; i++){
                turn_on(tour[i]);
            }
        }
        turn_on(v);
        ans += sum01;
        if (top){
            for (int i = down[v]; i < down[v]+sub[v]; i++){
                turn_off(tour[i]);
            }
        }
    };
    dfs(dfs,0,-1,true);
    ans *= 2;
    cout << ans << endl;
}
0