#956001 (C++23) No.2654 [Cherry 6th Tune] Re: start! (Black Sheep)

提出ソース

結果

問題	No.2654 [Cherry 6th Tune] Re: start! (Black Sheep)
ユーザー	umimel
提出日時	2024-02-26 18:19:33
言語	C++23 (gcc 12.3.0 + boost 1.83.0)
結果	AC
実行時間	343 ms / 7,000 ms
コード長	5,742 bytes
コンパイル時間	4,020 ms
コンパイル使用メモリ	259,132 KB
実行使用メモリ	59,772 KB
最終ジャッジ日時	2024-02-26 18:19:48
合計ジャッジ時間	12,624 ms
ジャッジサーバーID （参考情報）	judge15 / judge10

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	2 ms 6,676 KB
testcase_01	AC	2 ms 6,676 KB
testcase_02	AC	2 ms 6,676 KB
testcase_03	AC	175 ms 18,144 KB
testcase_04	AC	157 ms 17,396 KB
testcase_05	AC	85 ms 11,904 KB
testcase_06	AC	24 ms 6,676 KB
testcase_07	AC	67 ms 10,496 KB
testcase_08	AC	94 ms 12,800 KB
testcase_09	AC	154 ms 18,452 KB
testcase_10	AC	24 ms 6,676 KB
testcase_11	AC	101 ms 14,080 KB
testcase_12	AC	18 ms 6,676 KB
testcase_13	AC	98 ms 13,440 KB
testcase_14	AC	216 ms 22,596 KB
testcase_15	AC	194 ms 21,612 KB
testcase_16	AC	98 ms 11,904 KB
testcase_17	AC	171 ms 19,224 KB
testcase_18	AC	57 ms 9,728 KB
testcase_19	AC	226 ms 22,028 KB
testcase_20	AC	96 ms 13,440 KB
testcase_21	AC	110 ms 14,464 KB
testcase_22	AC	103 ms 14,208 KB
testcase_23	AC	2 ms 6,676 KB
testcase_24	AC	2 ms 6,676 KB
testcase_25	AC	2 ms 6,676 KB
testcase_26	AC	2 ms 6,676 KB
testcase_27	AC	3 ms 6,676 KB
testcase_28	AC	225 ms 23,676 KB
testcase_29	AC	270 ms 23,676 KB
testcase_30	AC	253 ms 23,676 KB
testcase_31	AC	242 ms 23,676 KB
testcase_32	AC	261 ms 23,548 KB
testcase_33	AC	343 ms 59,772 KB
testcase_34	AC	96 ms 22,248 KB
testcase_35	AC	236 ms 23,676 KB
testcase_36	AC	220 ms 59,772 KB
testcase_37	AC	246 ms 23,676 KB
testcase_38	AC	238 ms 59,772 KB
testcase_39	AC	2 ms 6,676 KB

権限があれば一括ダウンロードができます

ソースコード

raw source code

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
using pll = pair<ll, ll>;
#define drep(i, cc, n) for (ll i = (cc); i <= (n); ++i)
#define rep(i, n) drep(i, 0, n - 1)
#define all(a) (a).begin(), (a).end()
#define pb push_back
#define fi first
#define se second
mt19937_64 rng(chrono::system_clock::now().time_since_epoch().count());
const ll MOD1000000007 = 1000000007;
const ll MOD998244353 = 998244353;
const ll MOD[3] = {999727999, 1070777777, 1000000007};
const ll LINF = 1LL << 60LL;
const int IINF = (1 << 30) - 1;


template<typename T>
struct edge{
    int from, to;
    T cost;
    int id;

    edge(){}
    edge(int to, T cost=1) : from(-1), to(to), cost(cost){}
    edge(int to, T cost, int id) : from(-1), to(to), cost(cost), id(id){}
    edge(int from, int to, T cost, int id) : from(from), to(to), cost(cost), id(id){}
};

template<typename T>
struct redge{
    int from, to;
    T cap, cost;
    int rev;
    
    redge(int to, T cap, T cost=(T)(1)) : from(-1), to(to), cap(cap), cost(cost){}
    redge(int to, T cap, T cost, int rev) : from(-1), to(to), cap(cap), cost(cost), rev(rev){}
};

template<typename T> using Edges = vector<edge<T>>;
template<typename T> using weighted_graph = vector<Edges<T>>;
template<typename T> using tree = vector<Edges<T>>;
using unweighted_graph = vector<vector<int>>;
template<typename T> using residual_graph = vector<vector<redge<T>>>;


void solve(){
    int n; cin >> n;
    vector<ll> a(n+1);
    for(int i=0; i<n+1; i++) cin >> a[i];
    tree<int> T(n+1);
    for(int i=0; i<n; i++){
        int u, v; cin >> u >> v;
        T[u].push_back(edge<int>(v));
        T[v].push_back(edge<int>(u));
    }
    
    // lo <= hi <= lo+1
    multiset<ll> lo, hi;
    ll lsum=0, hsum=0;
    lo.insert(-LINF);
    hi.insert(LINF);
    auto add = [&](ll x){
        if((int)lo.size() == (int)hi.size()){
            auto itr = lo.end();
            itr--;
            if(*(itr) <= x){
                hi.insert(x);
                hsum += x;
            }else{
                hi.insert(*itr);
                hsum += *itr;
                lsum -= *itr;
                lo.erase(itr);
                lo.insert(x);
                lsum += x;
            }
        }else if((int)lo.size() < (int)hi.size()){
            if(*(hi.begin()) < x){
                lo.insert(*(hi.begin()));
                lsum += *(hi.begin());
                hsum -= *(hi.begin());
                hi.erase(hi.begin());
                hi.insert(x);
                hsum += x;
            }else{
                lo.insert(x);
                lsum += x;
            }
        }
    };

    auto erase = [&](ll x){
        if(lo.find(x)!=lo.end()){
            lo.erase(lo.lower_bound(x));
            lsum -= x;
            if((int)lo.size()+1 < (int)hi.size()){
                lo.insert(*(hi.begin()));
                lsum += *(hi.begin());
                hsum -= *(hi.begin());
                hi.erase(hi.begin());
            }
        }else{
            hi.erase(hi.lower_bound(x));
            hsum -= x;
            if((int)lo.size() > (int)hi.size()){
                auto itr = lo.end();
                itr--;
                hi.insert(*itr);
                hsum += *itr;
                lsum -= *itr;
                lo.erase(itr);
            }
        }
    };

    auto query = [&](){
        ll ans = LINF;
        //右端を選択
        {
            ll cost = 0;
            if((int)lo.size() == (int)hi.size()){
                // 中央値はunique
                auto itr = hi.end();
                itr--; itr--;
                ll mx = *itr;
                itr = lo.end();
                itr--;
                ll mid = *itr;
                cost += mid*(ll)(lo.size()-1) - lsum;
                cost += (hsum-mx) - mid*(ll)(hi.size()-2);

                if(mx == mid) cost++;
            }else{
                auto itr = hi.end();
                itr--; itr--;
                ll mx = *itr;
                itr = lo.end();
                itr--;
                ll lmid = *itr;
                ll rmid = *(hi.begin());
                cost += lmid*(ll)(lo.size()-1) - lsum;
                cost += (hsum-mx) - lmid*(ll)(hi.size()-2);
                if(lmid==rmid&&lmid==mx) cost++;
            }

            ans = min(ans, cost);
        }

        //左端を選択
        {
            ll cost = 0;
            if((int)lo.size() == (int)hi.size()){
                // 中央値はunique
                auto itr = lo.begin();
                itr++;
                ll mn = *itr;
                ll mid = *hi.begin();
                cost += mid*(ll)(lo.size()-2) - (lsum - mn);
                cost += hsum - mid*(ll)(hi.size()-1);

                if(mn == mid) cost++;
            }else{
                auto itr = lo.begin();
                itr++;
                ll mn = *itr;
                itr = hi.begin();
                ll lmid = *itr;
                itr++;
                ll rmid = *itr;
                cost += lmid*(ll)(lo.size()-2) - (lsum - mn);
                cost += hsum - lmid*(ll)(hi.size()-1);
                if(lmid==rmid&&lmid==mn) cost++;
            }

            ans = min(ans, cost);
        }

        return ans;
    };

    vector<ll> ans(n+1, -1), depth(n+1, 0);
    function<void(int, int)> dfs = [&](int v, int p){
        add(a[v]);
        if(depth[v]>=2) ans[v] = query();

        for(edge<int> e : T[v]) if(e.to!=p){
            depth[e.to] = depth[v]+1;
            dfs(e.to, v);
        }
        erase(a[v]);
    }; dfs(0, -1);

    for(int v=1; v<n+1; v++) cout << ans[v] << '\n';
}

int main(){
    cin.tie(nullptr);
    ios::sync_with_stdio(false);
    
    int T=1;
    // cin >> T;
    while(T--) solve();
}

yukicoder

結果

テストケース

ソースコード