Cod sursa(job #3146600)

Utilizator Razvan48Capatina Razvan Nicolae Razvan48 Data 21 august 2023 20:22:24
Problema Heavy Path Decomposition Scor 80
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 7.42 kb
#include <fstream>
#include <vector>

using namespace std;

const int NMAX = 1e5;
int val[1 + NMAX];

vector<int> graf[1 + NMAX];
bool vizitat[1 + NMAX];
int adancime[1 + NMAX];
vector<int> liniarizare;
int primaAparitie[1 + NMAX];

const int LMAX = 2 * NMAX - 1;

int aintLCA[1 + 4 * LMAX];

int heaviness[1 + NMAX]; /// ori dimensiunea subarborelui, ori lungimea maxima a unui lant ce pleaca din radacina subarborelui in jos

int indexPath[1 + NMAX];
int indexPePath[1 + NMAX];

struct Path
{
private:

    void buildAint(int index, int st, int dr)
    {
        if (st == dr)
        {
            this->aint[index] = val[this->noduri[st]];
            return;
        }

        int mij = (st + dr) / 2;
        int fiuSt = 2 * index;
        int fiuDr = 2 * index + 1;

        this->buildAint(fiuSt, st, mij);
        this->buildAint(fiuDr, mij + 1, dr);

        this->aint[index] = max(this->aint[fiuSt], this->aint[fiuDr]);
    }

    void updateAint(int index, int st, int dr, int poz, int val)
    {
        if (st == dr)
        {
            this->aint[index] = val;
            return;
        }

        int mij = (st + dr) / 2;
        int fiuSt = 2 * index;
        int fiuDr = 2 * index + 1;

        if (poz <= mij)
            this->updateAint(fiuSt, st, mij, poz, val);
        else
            this->updateAint(fiuDr, mij + 1, dr, poz, val);

        this->aint[index] = max(this->aint[fiuSt], this->aint[fiuDr]);
    }

    int queryAint(int index, int st, int dr, int stQ, int drQ)
    {
        if (st == stQ && drQ == dr)
            return this->aint[index];

        int mij = (st + dr) / 2;
        int fiuSt = 2 * index;
        int fiuDr = 2 * index + 1;

        if (drQ <= mij)
            return this->queryAint(fiuSt, st, mij, stQ, drQ);
        else if (stQ >= mij + 1)
            return this->queryAint(fiuDr, mij + 1, dr, stQ, drQ);
        else
            return max(this->queryAint(fiuSt, st, mij, stQ, mij), this->queryAint(fiuDr, mij + 1, dr, mij + 1, drQ));
    }

public:

    vector<int> noduri;
    int indexTataPath;
    int indexTataNodPath;

    vector<int> aint;

    Path()
    {
        this->noduri.emplace_back(-1);
    }

    Path(int indexTataPath, int indexTataNodPath) :
        indexTataPath(indexTataPath), indexTataNodPath(indexTataNodPath)
    {
        this->noduri.emplace_back(-1);
    }

    void buildAint()
    {
        this->aint.resize(4 * (int)this->noduri.size());

        this->buildAint(1, 1, (int)this->noduri.size() - 1);
    }

    void updateAint(int nod, int val)
    {
        this->updateAint(1, 1, (int)this->noduri.size() - 1, indexPePath[nod], val);
    }

    int queryAint(int nod1, int nod2)
    {
        if (indexPePath[nod1] > indexPePath[nod2])
            swap(nod1, nod2);

        return this->queryAint(1, 1, (int)this->noduri.size() - 1, indexPePath[nod1], indexPePath[nod2]);
    }
};

vector<Path> paths;

void dfs(int nod, int adanc)
{
    vizitat[nod] = true;
    adancime[nod] = adanc;
    liniarizare.emplace_back(nod);
    primaAparitie[nod] = (int)liniarizare.size() - 1;

    for (int i = 0; i < graf[nod].size(); ++i)
    {
        if (!vizitat[graf[nod][i]])
        {
            dfs(graf[nod][i], adanc + 1);
            liniarizare.emplace_back(nod);

            heaviness[nod] += heaviness[graf[nod][i]]; ///
        }
    }

    ++heaviness[nod]; ///
}

void buildAintLCA(int index, int st, int dr)
{
    if (st == dr)
    {
        aintLCA[index] = liniarizare[st];
        return;
    }

    int mij = (st + dr) / 2;
    int fiuSt = 2 * index;
    int fiuDr = 2 * index + 1;

    buildAintLCA(fiuSt, st, mij);
    buildAintLCA(fiuDr, mij + 1, dr);

    if (adancime[aintLCA[fiuSt]] < adancime[aintLCA[fiuDr]])
        aintLCA[index] = aintLCA[fiuSt];
    else
        aintLCA[index] = aintLCA[fiuDr];
}

int queryAintLCA(int index ,int st, int dr, int stQ, int drQ)
{
    if (st == stQ && drQ == dr)
        return aintLCA[index];

    int mij = (st + dr) / 2;
    int fiuSt = 2 * index;
    int fiuDr = 2 * index + 1;

    if (drQ <= mij)
        return queryAintLCA(fiuSt, st, mij, stQ, drQ);
    else if (stQ >= mij + 1)
        return queryAintLCA(fiuDr, mij + 1, dr, stQ, drQ);
    else
    {
        int sol1 = queryAintLCA(fiuSt, st, mij, stQ, mij);
        int sol2 = queryAintLCA(fiuDr, mij + 1, dr, mij + 1, drQ);

        if (adancime[sol1] < adancime[sol2])
            return sol1;
        else
            return sol2;
    }
}

int LCA(int nod1, int nod2)
{
    if (primaAparitie[nod1] > primaAparitie[nod2])
        swap(nod1, nod2);

    return queryAintLCA(1, 1, (int)liniarizare.size() - 1, primaAparitie[nod1], primaAparitie[nod2]);
}

void dfsPath(int nod)
{
    vizitat[nod] = true;
    paths.back().noduri.emplace_back(nod);
    indexPath[nod] = (int)paths.size() - 1;
    indexPePath[nod] = (int)paths.back().noduri.size() - 1;

    int maxHeaviness = 0;
    int indexFiuMaxHeaviness = -1;

    for (int i = 0; i < graf[nod].size(); ++i)
    {
        if (!vizitat[graf[nod][i]])
        {
            if (maxHeaviness < heaviness[graf[nod][i]])
            {
                maxHeaviness = heaviness[graf[nod][i]];
                indexFiuMaxHeaviness = i;
            }
        }
    }

    if (indexFiuMaxHeaviness != -1)
    {
        swap(graf[nod][0], graf[nod][indexFiuMaxHeaviness]);
        dfsPath(graf[nod][0]);
    }

    for (int i = 1; i < graf[nod].size(); ++i)
    {
        if (!vizitat[graf[nod][i]])
        {
            paths.emplace_back(indexPath[nod], nod);
            dfsPath(graf[nod][i]);
        }
    }
}

int solutiePeLant(int x, int y) /// y e undeva mai sus de x
{
    int sol = val[y];

    while (x != y)
    {
        if (indexPath[x] == indexPath[y])
        {
            sol = max(sol, paths[indexPath[x]].queryAint(x, y));
            x = y;
        }
        else
        {
            sol = max(sol, paths[indexPath[x]].queryAint(x, paths[indexPath[x]].noduri[1]));
            x = paths[indexPath[x]].indexTataNodPath;
        }
    }

    return sol;
}

int main()
{
    ifstream in("heavypath.in");
    ofstream out("heavypath.out");

    ios_base::sync_with_stdio(false);
    in.tie(nullptr);

    int n, m;
    in >> n >> m;

    for (int i = 1; i <= n; ++i)
        in >> val[i];

    for (int i = 2; i <= n; ++i)
    {
        int a, b;
        in >> a >> b;

        graf[a].emplace_back(b);
        graf[b].emplace_back(a);
    }

    liniarizare.reserve(1 + LMAX);

    liniarizare.emplace_back(0);
    dfs(1, 1);
    buildAintLCA(1, 1, (int)liniarizare.size() - 1);

    for (int i = 1; i <= n; ++i) /// nu am mai luat alt vector vizitat si pentru constructia path-urilor
        vizitat[i] = false;

    paths.emplace_back();
    paths.emplace_back(0, 0);
    dfsPath(1);

    for (int i = 1; i < paths.size(); ++i)
        paths[i].buildAint();

    for (int i = 1; i <= m; ++i)
    {
        int tip, x, y;
        in >> tip >> x >> y;

        if (tip == 0) /// update
        {
            val[x] = y;
            paths[indexPath[x]].updateAint(x, y);
        }
        else /// query
        {
            int lca = LCA(x, y);

            out << max(solutiePeLant(x, lca), solutiePeLant(y, lca)) << '\n';
        }
    }

    in.close();
    out.close();

    return 0;
}