Cod sursa(job #2763588)

Utilizator borcanirobertBorcani Robert borcanirobert Data 15 iulie 2021 13:15:10
Problema Heavy Path Decomposition Scor 60
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 9.42 kb
#include <iostream>
#include <algorithm>
#include <fstream>
#include <vector>
using namespace std;

class SegmentTree{
public:
    /// Creeaza un nou arbore de intervale pentru intervalul [x, y]
    /// Parametrii:
    ///      x: Inceputul intervalului
    ///      y: Finalul intervalului
    /// Complexitate: Theta(n*log(n))
    /// Complexitate spatiu: Theta(n*log(n)), unde n = y - x + 1
    SegmentTree(int x, int y):
        x{x}, y{y}, val{INITIAL_VALUE}, left{nullptr}, right{nullptr}
    {
        if (x != y)
        {
            int mid = (x + y) / 2;
            left = new SegmentTree(x, mid);
            right = new SegmentTree(mid + 1, y);
        }
    }

    SegmentTree(const SegmentTree& other)
        : x{other.x}, y{other.y}, val{other.val}
    {
        left = right = nullptr;
        if (x != y)
        {
            left = new SegmentTree(*(other.left));
            right = new SegmentTree(*(other.right));
        }
    }

    SegmentTree(SegmentTree&& other)
        :  x{other.x}, y{other.y}, val{other.val}, left{other.left}, right{other.right}
    {
        other.left = other.right = nullptr;
    }

    SegmentTree& operator=(const SegmentTree& other)
    {
        if (this == &other)
            return *this;

        x = other.x, y = other.y;
        val = other.val;

        delete left;
        delete right;

        left = right = nullptr;
        if (x != y)
        {
            left = new SegmentTree(*(other.left));
            right = new SegmentTree(*(other.right));
        }

        return *this;
    }

    SegmentTree& operator=(SegmentTree&& other)
    {
        if (this == &other)
            return *this;
        
        x = other.x, y = other.y;
        val = other.val;

        delete left;
        delete right;

        left = other.left;
        right = other.right;

        other.left = other.right = nullptr;

        return *this;
    }

    /// Actualizeaza valoarea de pe pozitia position cu valoarea value
    /// Parametrii:
    ///     position: Pozitia la care se actualizeaza valoarea
    ///     value: Valoarea cu care se actualizeaza
    /// Complexitate: O(log(n))
    void Update(int position, int value)
    {
        if (position < x || position > y)   // asta inseamna ca sunt intr-un interval care nu contine pozitia pe care o caut
            return;
        
        if (x == y && x == position)    // Am ajuns la frunza care retine valoarea de pe pozitia position
        {
            val = value;
            return;
        }

        // Trimit Update-ul spre stanga si spre dreapta, dar doar unul dintre fii va contine pozitia position
        left->Update(position, value);
        right->Update(position, value);

        // Actualizez valoarea din nodul curent, care retine informatii pentru tot intervalul [x, y]
        val = std::max(left->val, right->val);   // pentru maxime
    }

    /// Se gaseste suma/maximul/minimul pe intervalul [lo, hi]
    /// Parametrii:
    ///     lo: Inceputul intervalului pentru query
    ///     hi: Finalul intervalului pentru query
    /// Complexitate: O(log(n))
    int Query(int lo, int hi)   // Care este suma/minimul/maximul din intervalul [lo, hi]
    {
        if (y < lo || x > hi)   // sunt in afara intervalului pentru query
            return INITIAL_VALUE;
        
        if (lo <= x && y <= hi) // intervalul curent este inclus de intervalul pe care il caut
            return val;
        
        int t1 = left->Query(lo, hi);
        int t2 = right->Query(lo, hi);

        return std::max(t1, t2);
    }

    ~SegmentTree()
    {
        delete left;
        delete right;
    }
private:
    int x, y;   // intervalul pe care il retine nodul curent
    int val;    // valoarea de la acest nod - poate fi o suma, un minim, un maxim etc.
    SegmentTree *left, *right;      // fiul stang si fiul drept al arborelui

    const static int INITIAL_VALUE = 0;     // Ar trebui sa fie valoarea initiala din sir. Pentru sume e 0
};

class HeavyPath{
public:
    // Atentie: totul e indexat de la 1!!
    HeavyPath(int N, const std::vector<int>& values)
        : N{N}, values{values}, G{std::vector<std::vector<int>>(N + 1)}
    {}

    /// Adauga muchia [x, y] in arbore
    void addEdge(int x, int y)
    {
        G[x].push_back(y);
        G[y].push_back(x);
    }

    /// !!! Apeleaza aceasta functie inainte de query/update (dupa ce s-a construit tot arborele)
    void decomposeTree()
    {
        parent = lv = cnt = std::vector<int>(N + 1);
        where = std::vector<int>(N + 1);
        setsOfNodes.clear();
        DFS(1, 0);

        segmentTrees.clear();
        posInTree = std::vector<int>(N + 1);
        createTrees();

        // for (const auto& set : setsOfNodes)
        // {
        //     for (int x : set)
        //         std::cout << x << ' ' << where[x] << ' ' << posInTree[x] << '\n';
        //     std::cout << '\n';
        // }
        // std::cout << segmentTrees[0].Query(2, 3) << '\n';
        // std::cout << segmentTrees[0].Query(1, 3) << '\n';
        // std::cout << segmentTrees[0].Query(3, 4) << '\n';
        // std::cout << segmentTrees[3].Query(2, 3) << '\n';
    }

    /// Actualizez valoarea din nodul node cu noua valoare value
    /// Parametrii:
    ///     node: nodul in care se actulizeaza valoarea
    ///     value: noua valoare pe care o va avea node
    void update(int node, int value)
    {
        int index = where[node];
        int indexInTree = posInTree[node];
        segmentTrees[index].Update(indexInTree, value);
        values[node] = value;
    }

    /// Query pentru a afla maximul pe lantul [x, y], unde x e stramos al lui y sau y e stramos al lui x
    /// Parametrii:
    ///     x: nod din arbore
    ///     y: nod din arbore
    /// Returneaza maximul dintre valorile de pe lantul [x, y]
    int query(int x, int y)
    {
        if (where[x] == where[y])
        {
            if (lv[x] < lv[y])
               std::swap(x, y);

            int posX = posInTree[x];
            int posY = posInTree[y];
            return segmentTrees[where[x]].Query(posX, posY);
        }

        if (lv[setsOfNodes[where[x]].back()] < lv[setsOfNodes[where[y]].back()])
            std::swap(x, y);

        int posX = posInTree[x];
        int pFinal = static_cast<int>(setsOfNodes[where[x]].size());
        int ancestorX = parent[setsOfNodes[where[x]].back()];

        int currentResult = segmentTrees[where[x]].Query(posX, pFinal);
        return std::max(currentResult, query(ancestorX, y));
    }

private:
    void DFS(int node, int parent)
    {
        cnt[node] = 1;
        for (int desc : G[node])
            if (desc != parent)
            {
                this->parent[desc] = node;
                lv[desc] = lv[node] + 1;
                DFS(desc, node);
                cnt[node] += cnt[desc];
            }
        
        if (cnt[node] == 1)     // node e frunza
        {
            setsOfNodes.push_back(std::vector<int>(1, node));
            where[node] = static_cast<int>(setsOfNodes.size()) - 1;
        }
        else
        {
            int bestDesc{-1};
            for (int desc : G[node])
                if (desc != parent)
                {
                    if (bestDesc == -1 || cnt[desc] > cnt[bestDesc])
                        bestDesc = desc;
                }
            
            setsOfNodes[where[bestDesc]].push_back(node);
            where[node] = where[bestDesc];
        }
    }

    void createTrees()
    {
        for (const auto& currentSet : setsOfNodes)
        {
            segmentTrees.emplace_back(1, static_cast<int>(currentSet.size()));
            for (size_t i = 0; i < currentSet.size(); ++i)
            {
                segmentTrees.back().Update(i + 1, values[currentSet[i]]);
                posInTree[currentSet[i]] = i + 1;
            }
        }
    }

    int N;
    std::vector<int> values;
    std::vector<std::vector<int>> G;
    std::vector<int> lv, cnt;                       // lv[i] = nivelul nodului i, cnt[i] = nr. de noduri din subarborele i
    std::vector<int> parent;                        // parent[i] = parintele nodului i
    std::vector<std::vector<int>> setsOfNodes;      // setsOfNodes = lista cu seturile de noduri ce vor forma arbori de intervale
    std::vector<int> where;                         // where[i] = in care arbore se afla nodul i
    std::vector<int> posInTree;                     // posInTree[i] = pozitia in arborele where[i] la care am nodul i
    std::vector<SegmentTree> segmentTrees;          // segmentTrees = vector cu toti arborii de intervale
};

std::ifstream fin("heavypath.in");

HeavyPath* H;
int M;

void ReadData();
void Solve();

int main()
{
    ReadData();
    Solve();

    return 0;
}

void ReadData()
{
    int N, x, y;
    fin >> N >> M;

    std::vector<int> a(N + 1);
    for (int i = 1; i <= N; ++i)
        fin >> a[i];

    H = new HeavyPath(N, a);

    for (int i = 1; i < N; ++i)
    {
        fin >> x >> y;

        H->addEdge(x, y);
    }
}

void Solve()
{
    std::ofstream fout("heavypath.out");
    int type, x, y;

    H->decomposeTree();
    while (M--)
    {
        fin >> type >> x >> y;

        if (type == 0)
            H->update(x, y);
        else
            fout << H->query(x, y) << '\n';
    }

    delete H;
    fout.close();
}