Submission #1030931

#	Time	Username	Problem	Language	Result	Execution time	Memory
1030931		Antekb	JOI tour (JOI24_joitour)	C++17	0 / 100	3035 ms	418548 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 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;
        ll 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[i]={a.st, b.st};
            sum[i]={a.nd, 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;
    }
    drzewo(vi V, vector<pii> 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){
            u=new_name(u);
            v=new_name(v);
            E[u].pb(v);
            E[v].pb(u);
        }
        dfs_cent(0);
        assert(root!=-1);
        real_root=real_name[root];
        par[root]=0;
        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();
    }
    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();
    }
};

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) {
    typ[X]=Y;
    while(X!=-1){
        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:104:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  104 |         for(int i=0; i<E[root].size(); i++){
      |                      ~^~~~~~~~~~~~~~~

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