Cod sursa(job #3146343)

Utilizator Razvan48Capatina Razvan Nicolae Razvan48 Data 20 august 2023 15:22:22
Problema Heavy Path Decomposition Scor 40
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 7.14 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 log2[1 + LMAX];

const int LOGMAX = 17;
int rmq[1 + LOGMAX][1 + 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 preCalcRMQ()
{
    log2[0] = 0;
    log2[1] = 0;
    for (int i = 2; i <= LMAX; ++i)
        log2[i] = log2[i / 2] + 1;

    for (int i = 1; i < liniarizare.size(); ++i)
        rmq[0][i] = liniarizare[i];

    for (int log = 1; log <= LOGMAX; ++log)
    {
        for (int i = 1; i + (1 << log) - 1 < liniarizare.size(); ++i)
        {
            if (adancime[rmq[log - 1][i]] < adancime[rmq[log - 1][i + (1 << (log - 1))]])
                rmq[log][i] = rmq[log - 1][i];
            else
                rmq[log][i] = rmq[log - 1][i + (1 << (log - 1))];
        }
    }
}

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

    int lungime = primaAparitie[nod2] - primaAparitie[nod1] + 1;
    int logLungime = log2[lungime];

    if (adancime[rmq[logLungime][primaAparitie[nod1]]] < adancime[rmq[logLungime][primaAparitie[nod2] - (1 << logLungime) + 1]])
        return rmq[logLungime][primaAparitie[nod1]];
    else
        return rmq[logLungime][primaAparitie[nod2] - (1 << logLungime) + 1];
}

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.emplace_back(0);
    dfs(1, 1);
    preCalcRMQ();

    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;
}