Submission #782339

#	Time	Username	Problem	Language	Result	Execution time	Memory
782339		anton	Construction of Highway (JOI18_construction)	C++17	100 / 100	991 ms	102780 KiB

construction

#include<bits/stdc++.h>

using namespace std;
#define int long long
#define pii pair<int, int>
vector<int> colors;
vector<int> time_g;
vector<int> depth;

map<int, int> m;
struct Node{
    int t;
    int last_m;
    int color;
    mutable vector<pii> cols;

    bool operator<( const Node& b)const{
        return t<b.t;
    }

    Node(int id){
        t = time_g[id];
        color = colors[id];
        last_m =id;
        //cols.push_back(pii(depth[id]-1, color));
    }
};
vector<vector<pii>> adj;
vector<set<Node>> sets;

void merge_pure(set<Node>&s, set<Node>&b, int prov){
    for(auto e: b){
        auto e2 = e;
        e2.last_m = prov;
        s.insert(e);
    }
}

void compute_rel(set<Node>::iterator a, set<Node>::iterator b, int d){
    //cout<<a->t<<" "<<b->t<<endl;
    if(a->last_m!=b->last_m){
        //cout<<m[a->t]+1<<" includes "<<m[b->t]+1<<endl;
        a->cols.push_back(pii(d, b->color));
    }
}

void merge(set<Node>&a, set<Node>& b, int id){
    for(auto it = b.begin(); it!=b.end(); it++){
        auto pre = a.lower_bound((*it));

        if(pre!=a.end()){
            it++;
            if(it== b.end() || pre->t<it->t){
                it--;
                compute_rel(pre, it, depth[id]);
                it++;
            }
            it--;
        }
        if(pre!=a.begin()){
            pre--;
            if(pre!=a.end() && ((*pre)<(*it))){
                compute_rel(it, pre, depth[id]);
                
            }
        }
        a.insert(*it);
    }

}
void dfs(int id){
    set<Node> res;
    
    set<Node> comp;
    if(adj[id].size()>0){
        pii best = pii(0, 0);
        for(auto ch: adj[id]){
            dfs(ch.first);
            if(sets[ch.first].size()>best.first){
                best.first = sets[ch.first].size();
                best.second = ch.first;
            }
        }
        swap(res, sets[best.second]);
        for(auto e: adj[id]){
            if(e.first != best.second){
                merge_pure(comp, sets[e.first],e.first);
                sets[e.first].clear();
            }
        }

    }

        
    comp.insert(Node(id));
    merge(res, comp, id);
    swap(res, sets[id]);
}

int merge(vector<pii>&v, int lt, int rt, int mid){
    
    vector<pii> res;
    int c= 0;
    int l = lt;
    int r= mid+1;

    int r_s =0;
    for(int i = lt; i<=rt; i++){
        if(l>mid){
            res.push_back(v[r]);
            r_s += v[r].first;
            r++;
        }
        else if(r>rt){
            res.push_back(v[l]);
            c+=(r_s)*v[l].first;
            l++;
        }
        else if(v[l].second<=v[r].second){
            res.push_back(v[l]);
            c+=(r_s)*v[l].first;
            l++;
        }
        else{
            res.push_back(v[r]);
            r_s += v[r].first;
            r++;
        }
    }

    for(int i = lt; i<=rt; i++){
        v[i] = res[i-lt];
    }
    return c;
}

int dpr(vector<pii>&v, int lt, int rt){
    if(lt==rt){
        return 0;
    }
    else{
        //cout<<lt<<" "<<rt<<endl;
        int mid = (lt+rt)/2;
        int r =0;
        r+=dpr(v, lt, mid);
        r+=dpr(v, mid+1, rt);
        r+= merge(v, lt, rt, mid);
        //cout<<r<<endl;
        return r;
    }
}

signed main(){
    int n;
    cin>>n;
    colors.resize(n);
    adj.resize(n);
    time_g.resize(n, -1);
    depth.resize(n, 1);
    sets.resize(n);

    for(int i = 0; i<n; i++){
        int c;
        cin>>c;
        colors[i] = c;
    }

    m.insert(pii(-1, 0));

    for(int i = 0; i<n-1; i++){
        int a, b;
        cin>>a>>b;
        a--;b--;
        time_g[b] = i;

        adj[a].push_back(pii(b, i));
        depth[b] = depth[a]+1;
        m.insert(pii(i, b));
    }

    dfs(0);
    set<Node> res;
    swap(res, sets[0]);

    vector<int> ans(n-1);

    for(auto e: res){
        int cur= e.t;
        if(cur!=-1){

            int r= 0;

            sort(e.cols.begin(), e.cols.end());

            for(int i = e.cols.size()-1; i>0; i--){
                e.cols[i].first -=e.cols[i-1].first;
            }
            r += dpr(e.cols, 0, e.cols.size()-1);
            ans[cur]= r;
        }
    }
    for(int i= 0; i<n-1; i++){
        cout<<ans[i]<<endl;
    }
}

Compilation message (stderr)

construction.cpp: In function 'void dfs(long long int)':
construction.cpp:79:37: warning: comparison of integer expressions of different signedness: 'std::set<Node>::size_type' {aka 'long unsigned int'} and 'long long int' [-Wsign-compare]
   79 |             if(sets[ch.first].size()>best.first){
      |                ~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~

• Subtask 1: 7 / 7
• Subtask 2: 9 / 9
• Subtask 3: 84 / 84