Cod sursa(job #2646769)

Utilizator stefan.popescuPopescu Stefan stefan.popescu Data 1 septembrie 2020 23:48:38
Problema Heavy Path Decomposition Scor 60
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 5.5 kb
#include <iostream>
#include <fstream>
#include <vector>
using namespace std;
ifstream in ("heavypath.in");
ofstream out("heavypath.out");
int n, qq, x, y, mom, tip;
vector <int> firstAp; ///prima aparitie in eulerTour
vector <vector <int> > vecini;
vector <int> visIn; ///momentul vizitei initiale
vector <int> order; ///ordinea in care apar elementele in parcurgerea dfs2
vector <int> depth; ///adancimile nodurilor
vector <int> vals; ///valorile nodurilor (initiale)
vector <int> seg; ///segtree-ul
vector <int> component; ///capul lantului
vector <int> eulerTour; ///parcurgerea Euler
vector <int> father; ///tatal fiecarui nod
int rmq[18][200000];
int loga[200001];
vector <int> subSum; ///numarul de elemente din subarbore

void getReadingDone();

void getHLDBuilt();
void dfs1(int nod, int dad);
void dfs2(int nod);
void buildSeg(int st=0, int dr=n-1, int poz=0);

void getLcaBuilt();
void eulerDfs(int nod, int level);
int getLca(int x, int y);

void update(int pozQQ, int val, int st=0, int dr=n-1, int poz=0);
int query(int stQQ, int drQQ, int st=0, int dr=n-1, int poz=0);

int queryPath(int x, int y){
    int lca=getLca(x, y);

    int maxX=0, maxY=0;

    if(component[x]==component[lca]){
        maxX=query(visIn[lca], visIn[x]);
    }
    else {
        int ant;
        do{
            ant=component[x];
            maxX=max(maxX, query(visIn[ant], visIn[x]) );
            x=father[ant];
        }while(component[x]!=component[lca]);
        maxX=max(maxX, query(visIn[lca], visIn[x]));
    }

    if(component[y]==component[lca]){
        maxY=query(visIn[lca], visIn[y]);
    }
    else {
        int ant;
        do{
            ant=component[y];
            maxY=max(maxY, query(visIn[ant], visIn[y]) );
            y=father[ant];
        }while(component[y]!=component[lca]);
        maxY=max(maxY, query(visIn[lca], visIn[y]));
    }

    return max(maxX, maxY);
}

int main()
{
    getReadingDone();

    getHLDBuilt();

    getLcaBuilt();

    while(qq--){
        in>>tip>>x>>y;
        if(tip==0){
            update(visIn[x], y);
        }
        else {
            out<<queryPath(x, y)<<"\n";
        }
    }
    return 0;
}

void getReadingDone(){
    in>>n>>qq;
    vals.resize(n+1);
    vecini.resize(n+1);
    for(int i=1; i<=n ; i++){
        in>>vals[i];
    }
    for(int i=1; i<n; i++){
        in>>x>>y;
        vecini[x].push_back(y);
        vecini[y].push_back(x);
    }
}


void getHLDBuilt(){
    subSum.resize(n+1);
    father.resize(n+1);
    dfs1(1, 0);

    mom=0;
    order.resize(n+1);
    visIn.resize(n+1);
    component.resize(n+1);
    component[1]=1;
    dfs2(1);

    seg.resize(2*n+1);
    buildSeg();
}

void dfs1(int nod, int dad){
    father[nod]=dad;
    subSum[nod]=1;
    for(unsigned i=0; i<vecini[nod].size(); i++)
        if(vecini[nod][i]!=dad){
            dfs1(vecini[nod][i], nod);

            subSum[nod]+=subSum[vecini[nod][i]];

            if(subSum[vecini[nod][i]]>subSum[vecini[nod][0]]||vecini[nod][0]==dad)
                swap(vecini[nod][i], vecini[nod][0]);
        }
}

void dfs2(int nod){
    visIn[nod]=mom;
    order[mom]=nod;
    mom++;
    for(unsigned i=0; i<vecini[nod].size(); i++)
        if(vecini[nod][i]!=father[nod]){
            component[vecini[nod][i]]=(vecini[nod][i]==vecini[nod][0])? component[nod] : vecini[nod][i];
            dfs2(vecini[nod][i]);
        }
}

void buildSeg(int st, int dr, int poz){
    if(st==dr){
        seg[poz]=vals[order[st]];
        return;
    }
    int mij=(st+dr)/2;
    buildSeg(st, mij, poz+1);
    buildSeg(mij+1, dr, poz+2*(mij-st+1));
    seg[poz]=max(seg[poz+1], seg[poz+2*(mij-st+1)]);
}


void getLcaBuilt(){
    loga[1]=0;
    for(int i=2; i<=2*n; i++)
        loga[i]=loga[i>>1]+1;

    depth.resize(n+1);
    firstAp.resize(n+1);
    mom=0;
    eulerDfs(1, 1);

    for(int put=1; (1<<put)<=mom; put++)
        for(int i=0; i+(1<<put)-1<mom; i++){
            if(depth[rmq[put-1][i]] < depth[rmq[put-1][i+(1<<(put-1))]])
                rmq[put][i]=rmq[put-1][i];
            else
                rmq[put][i]=rmq[put-1][i+(1<<(put-1))];
        }
}

void eulerDfs(int nod, int level){
    depth[nod]=level;
    rmq[0][mom++]=nod;
    firstAp[nod]=mom-1;
    for(auto &x:vecini[nod])
        if(father[nod]!=x){
            eulerDfs(x, level+1);
            rmq[0][mom++]=nod;
        }
}


int getLca(int x, int y){
    if(x==y) return x;
    if(firstAp[x]>firstAp[y])
        swap(x, y);

    int p1=firstAp[x], p2=firstAp[y];
    int putt=loga[p2-p1+1];

    if(depth[rmq[putt][p1] ]<depth[rmq[putt][p2-(1<<putt)+1]] )
        return rmq[putt][p1];
    else
        return rmq[putt][p2-(1<<putt)+1];
}

void update(int pozQQ, int val, int st, int dr, int poz){
    if(st==dr){
        seg[poz]=val;
        return;
    }
    int mij=(st+dr)/2;
    if(pozQQ<=mij)
        update(pozQQ, val, st, mij, poz+1);
    else
        update(pozQQ, val, mij+1, dr, poz+2*(mij-st+1));
    seg[poz]=max(seg[poz+1], seg[poz+2*(mij-st+1)]);
}

int query(int stQQ, int drQQ, int st, int dr, int poz){
    if(stQQ<=st&&dr<=drQQ)
        return seg[poz];
    int mij=(st+dr)/2;
    if(stQQ<=mij&&mij+1<=drQQ)
        return max(query(stQQ, drQQ, st, mij, poz+1), query(stQQ, drQQ, mij+1, dr, poz+2*(mij-st+1)) );
    else if (stQQ<=mij)
        return query(stQQ, drQQ, st, mij, poz+1);
    else if(mij+1<=drQQ)
        return query(stQQ, drQQ, mij+1, dr, poz+2*(mij-st+1));
}