Submission #1031071

#	Time	Username	Problem	Language	Result	Execution time	Memory
1031071		Antekb	JOI tour (JOI24_joitour)	C++17	100 / 100	2270 ms	522528 KiB

joitour

#include "bits/stdc++.h"	/** keep-include */
using namespace std;

#define rep(i, b, e) for(int i = (b); i <= (e); i++)
#define per(i, b, e) for(int i = (e); i >= (b); i--)
#define FOR(i, b, e) rep(i, b, (e) - 1)
#define SZ(x) int(x.size())
#define all(x) x.begin(), x.end()
#define pb push_back
#define mp make_pair
#define st first
#define nd second
using ll = long long;
using vi = vector<int>;
using pii = pair<int, int>;

auto &operator<<(auto &o, pair<auto, auto> p) {
	return o << "(" << p.st << ", " << p.nd << ")"; }
auto operator<<(auto &o, auto x)->decltype(end(x), o) {
	o << "{"; int i=0; for(auto e: x) o << ", " + 2*!i++ << e;
	return o << "}"; }
#ifdef LOCAL
#define deb(x...) cerr << "[" #x "]: ", [](auto...$) { \
	((cerr << $ << "; "),...) << endl; }(x)
#else
#define deb(...)
#endif

#include "joitour.h"

const int N=2e5+5;

ll result;
int typ[N];


struct drzewo{
    vi real_name;
    vi siz, par, pre, post, subtree, order;
    vector<vi> E;
    int root=-1, real_root;
    int n;
    vi ile[2];
    vector<ll> sum[2];
    ll iles[2]={}, sums[2]={};
    ll zle=0;
    ll res=0;
    int new_name(int v){
        return lower_bound(all(real_name), v)-real_name.begin();
    }
    void dfs_cent(int v){
        siz[v]=1;
        bool ok=1;
        for(int u:E[v]){
            if(u!=par[v]){
                par[u]=v;
                dfs_cent(u);
                siz[v]+=siz[u];
                if(siz[u]*2>n)ok=0;
            }
        }
        if(ok && 2*(n-siz[v])<=n){
            root=v;
        }
    }
    void dfs_init(int v, int sub=-1){
        deb(v, sub, order, root, E[v]);
        pre[v]=order.size();
        order.pb(v);
        int cnt=0;
        if(v!=root){
            subtree[v]=sub;
            for(int &u:E[v])if(u==par[v])swap(u, E[v].back());
            E[v].pop_back();
        }
        for(int u:E[v]){
            par[u]=v;
            if(v==root)dfs_init(u, cnt);
            else dfs_init(u, sub);
            cnt++;
        }
        post[v]=order.size();
    }
    pair<int, ll> licz(int v, int lbl, int num){//num - liczba 1
        pair<int, ll> a;
        for(int u:E[v]){
            auto b=licz(u, lbl, num+(typ[real_name[v]]==1));
            a.st+=b.st;
            a.nd+=b.nd;
        }
        if(typ[real_name[v]]==lbl*2){
            a.st++;
            a.nd+=num;
        }
        return a;
    }
    void calc(){
        result-=res;
        res=0;
        iles[0]=0;
        iles[1]=0;
        sums[0]=0;
        sums[1]=0;
        zle=0;
        for(int i=0; i<E[root].size(); i++){
            auto a=licz(E[root][i], 0, 0);
            auto b=licz(E[root][i], 1, 0);
            ile[0][i]={a.st};
            ile[1][i]={b.st};
            sum[0][i]={a.nd};
            sum[1][i]={b.nd};
            iles[0]+=a.st;
            iles[1]+=b.st;
            sums[0]+=a.nd;
            sums[1]+=b.nd;
            res-=a.st*b.nd;
            res-=a.nd*b.st;
            zle+=a.st*1ll*b.st;
        }
        res+=sums[0]*iles[1];
        res+=sums[1]*iles[0];
        if(typ[real_name[root]]==1)res+=iles[0]*1ll*iles[1]-zle;
        else res+=sums[1-typ[real_name[root]]/2];
        result+=res;
    }
    vector<int> suma[2];
    vector<int> srod;
    void add_srod(int l, int r, int c){
        for(l+=n, r+=n; l<r; l>>=1, r>>=1){
            if(l&1)srod[l++]+=c;
            if(r&1)srod[--r]+=c;
        }
    }
    int get_srod(int v){
        v+=n;
        int ans=0;
        while(v>0){
            ans+=srod[v];
            v/=2;
        }
        return ans;
    }
    int licz_suma(int l, int r, int ktora){
        int ans=0;
        for(l+=n, r+=n; l<r; l>>=1, r>>=1){
            if(l&1)ans+=suma[ktora][l++];
            if(r&1)ans+=suma[ktora][--r];
        }
        return ans;
    }
    void add_suma(int v, int c, int ktora){
        v+=n;
        while(v>0){
            suma[ktora][v]+=c;
            v/=2;
        }
    }
    drzewo(vi V, vector<pii> edges){
        deb(V, edges);
        real_name=V;
        sort(all(real_name));
        n=real_name.size();
        E.resize(n);
        siz.resize(n);
        par.resize(n);
        subtree.resize(n);
        pre.resize(n);
        post.resize(n);
        par[0]=-1;
        for(auto &[u, v] : edges){
            deb(u, v, real_name);
            u=new_name(u);
            v=new_name(v);
            E[u].pb(v);
            E[v].pb(u);
        }
        dfs_cent(0);
        assert(root!=-1);
        deb(root, real_name);
        real_root=real_name[root];
        par[root]=-1;
        ile[0].resize(E[root].size());
        ile[1].resize(E[root].size());
        sum[0].resize(E[root].size());
        sum[1].resize(E[root].size());
        dfs_init(root);
        calc();
        suma[0].resize(n+n);
        suma[1].resize(n+n);
        srod.resize(n+n);
        for(int i=0; i<n; i++){
            if(i==root)continue;
            if(typ[real_name[i]]==1){
                add_srod(pre[i], post[i], 1);
            }
            else{
                add_suma(pre[i], 1, typ[real_name[i]]/2);
            }
        }
    }
    pair<vi, vector<pii> > prep(int v){
        vi vert;
        vector<pii> edg;
        for(int i=pre[v]; i<post[v]; i++){
            int u=order[i];
            vert.pb(real_name[u]);
            for(int w:E[u]){
                edg.pb(mp(real_name[u], real_name[w]));
            }
        }
        return {vert, edg};
    }
    void update_slow(){
        calc();
        deb(res);
    }
    int licz_gora(int v){
        deb(v, typ[real_name[v]]);
        if(v==root)return 0;
        return licz_gora(par[v])+(typ[real_name[v]]==1);
    }
    int licz_gora2(int v){
        return get_srod(pre[v]);
    }
    int licz_dol(int v, int lbl){
        int cnt=0;
        for(int i=pre[v]; i<post[v]; i++){
            if(typ[real_name[order[i]]]==lbl)cnt++;
        }
        return cnt;
    }
    int licz_dol2(int v, int lbl){
        return licz_suma(pre[v], post[v], lbl/2);
    }

    void update_fast(int v, int stary, int nowy){
        v=new_name(v);
        deb(root, v, stary, nowy);
        deb(ile[0], ile[1], sum[0], sum[1]);
        result-=res;
        typ[real_name[v]]=stary;
        
        if(typ[real_name[root]]==1)res-=iles[0]*1ll*iles[1]-zle;
        else res-=sums[1-typ[real_name[root]]/2];
        
        if(v!=root){
            res-=sums[0]*iles[1];
            res-=sums[1]*iles[0];
            int sub=subtree[v];

            iles[0]-=ile[0][sub];
            iles[1]-=ile[1][sub];
            sums[0]-=sum[0][sub];
            sums[1]-=sum[1][sub];
            res+=ile[0][sub]*sum[1][sub];
            res+=ile[1][sub]*sum[0][sub];
            zle-=ile[1][sub]*1ll*ile[0][sub];

            if(stary==1){
                int cnt0=licz_dol2(v, 0);
                int cnt2=licz_dol2(v, 2);
                sum[0][sub]-=cnt0;
                sum[1][sub]-=cnt2;
                add_srod(pre[v], post[v], -1);
            }
            else{
                int cnt=licz_gora2(v);
                //deb(cnt, licz_gora2(v));
                //assert(cnt==licz_gora2(v));
                sum[stary/2][sub]-=cnt;
                ile[stary/2][sub]--;
                add_suma(pre[v], -1, typ[real_name[v]]/2);
            }
            deb(ile[0], ile[1], sum[0], sum[1]);
            deb(res);
            typ[real_name[v]]=nowy;
            
            if(nowy==1){
                int cnt0=licz_dol2(v, 0);
                int cnt2=licz_dol2(v, 2);
                sum[0][sub]+=cnt0;
                sum[1][sub]+=cnt2;
                add_srod(pre[v], post[v], 1);
            }
            else{
                int cnt=licz_gora2(v);
                //deb(cnt, licz_gora2(v));
                //assert(cnt==licz_gora2(v));
                sum[nowy/2][sub]+=cnt;
                ile[nowy/2][sub]++;
                add_suma(pre[v], 1, typ[real_name[v]]/2);
            }

            iles[0]+=ile[0][sub];
            iles[1]+=ile[1][sub];
            sums[0]+=sum[0][sub];
            sums[1]+=sum[1][sub];
            res-=ile[0][sub]*sum[1][sub];
            res-=ile[1][sub]*sum[0][sub];
            zle+=ile[1][sub]*1ll*ile[0][sub];
            
            res+=sums[0]*iles[1];
            res+=sums[1]*iles[0];
        }
        typ[real_name[v]]=nowy;
        if(typ[real_name[root]]==1)res+=iles[0]*1ll*iles[1]-zle;
        else res+=sums[1-typ[real_name[root]]/2];
        result+=res;
        deb(root, v, stary, nowy);
        deb(ile[0], ile[1], sum[0], sum[1]);
        deb(res);
    }
};

vector<drzewo*> drzewa(N);
int rodzic[N];

void stworz(vector<int> vert, vector<pii> edg, int par){
    drzewo* T=new drzewo(vert, edg);
    drzewa[T->real_root]=T;
    deb(T->real_root);
    rodzic[T->real_root]=par;
    for(int u:T->E[T->root]){
        auto a = T->prep(u);
        deb(T->real_root, a.st, a.nd);
        stworz(a.st, a.nd, T->real_root);
    }
}

int n;

void init(int _n, std::vector<int> F, std::vector<int> U, std::vector<int> V,
          int Q) {
    n=_n;
    vector<pii> edges;
    for(int i=0; i<n-1; i++){
        edges.pb(mp(U[i], V[i]));   
    }
    for(int i=0; i<n; i++){
        typ[i]=F[i];
    }
    vi v(n);
    iota(all(v), 0);
    stworz(v, edges, -1);
}

void change(int X, int Y) {
    int stary=typ[X];
    int v=X;
    typ[X]=Y;
    while(X!=-1){
        drzewa[X]->update_fast(v, stary, Y);
        //drzewa[X]->update_slow();
        X=rodzic[X];
    }
}

long long num_tours() {
  return result;
}

Compilation message (stderr)

joitour.cpp:17:18: warning: use of 'auto' in parameter declaration only available with '-fconcepts-ts'
   17 | auto &operator<<(auto &o, pair<auto, auto> p) {
      |                  ^~~~
joitour.cpp:17:32: warning: use of 'auto' in parameter declaration only available with '-fconcepts-ts'
   17 | auto &operator<<(auto &o, pair<auto, auto> p) {
      |                                ^~~~
joitour.cpp:17:38: warning: use of 'auto' in parameter declaration only available with '-fconcepts-ts'
   17 | auto &operator<<(auto &o, pair<auto, auto> p) {
      |                                      ^~~~
joitour.cpp:19:17: warning: use of 'auto' in parameter declaration only available with '-fconcepts-ts'
   19 | auto operator<<(auto &o, auto x)->decltype(end(x), o) {
      |                 ^~~~
joitour.cpp:19:26: warning: use of 'auto' in parameter declaration only available with '-fconcepts-ts'
   19 | auto operator<<(auto &o, auto x)->decltype(end(x), o) {
      |                          ^~~~
joitour.cpp: In member function 'void drzewo::calc()':
joitour.cpp:105:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  105 |         for(int i=0; i<E[root].size(); i++){
      |                      ~^~~~~~~~~~~~~~~

• Subtask 1: 6 / 6
• Subtask 2: 8 / 8
• Subtask 3: 6 / 6
• Subtask 4: 16 / 16
• Subtask 5: 16 / 16
• Subtask 6: 34 / 34
• Subtask 7: 14 / 14