Cod sursa(job #2025603)

Utilizator borcanirobertBorcani Robert borcanirobert Data 22 septembrie 2017 21:49:38
Problema Heavy Path Decomposition Scor 80
Compilator cpp Status done
Runda Arhiva educationala Marime 4.76 kb
#include <algorithm>
#include <iostream>
#include <fstream>
#include <vector>
using namespace std;

ifstream fin("heavypath.in");
ofstream fout("heavypath.out");

using VI = vector<int>;
using Graph = vector<VI>;
struct Tree{
public:
    Tree(int N1, int N2);
    void Update(int pos, int value);
    int Query(int p1, int p2);

    int i1, i2;
    int val;
    Tree *leftNode, *rightNode;
};
int N, M;
VI cnt, a, f, ind, lv, viz;
Graph G, G0;
vector<Tree> trees;
VI treeSize;
int cntTrees;
vector<vector<int>> v;

void Read();
void Decompose();
void Update(int node, int val);
int Query(int p1, int p2);
void InitialDFS( int node );
void DFS( int node );
int FindPos( int node );

bool CompLevels( const int& a, const int& b )
{
    return lv[a] < lv[b];
}

int main()
{
    Read();
    InitialDFS(1);
    Decompose();

   // cout << trees[3].val; cin.get();

    int type, x, y;
    while ( M-- )
    {
        fin >> type >> x >> y;

        if (type == 0)
            Update(x, y);
        else
            fout << Query(x, y) << '\n';
    }

    fin.close();
    fout.close();
    return 0;
}

void Read()
{
    fin >> N >> M;
    viz = lv = ind = cnt = f = a = VI(N + 1);
    f[1] = -1;
    G0 = G = Graph(N + 1);

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

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

        G0[x].push_back(y);
        G0[y].push_back(x);
    }
}

void InitialDFS( int node )
{
    viz[node] = true;
    for ( const int& son : G0[node] )
        if ( !viz[son] )
        {
            InitialDFS(son);
            f[son] = node;
            G[node].push_back(son);
        }
}

void Decompose()
{
    DFS(1);

    /*for ( int i = 1; i <= N; ++i )
        cout << i << " : " << lv[i] << '\n';
    cin.get();  */

   /* for ( int i = 1; i <= N; ++i )
        fout << i << " : " << ind[i] << '\n';   */
    for ( const int& x : treeSize )
    {
        Tree *t = new Tree(0, x - 1);
        trees.push_back(*t);
    }

    v = vector<vector<int>>(cntTrees);
    for ( int i = 1; i <= N; ++i )
        v[ind[i]].push_back(i);
    for ( int i = 0; i < cntTrees; ++i )
    {
        sort(v[i].begin(), v[i].end(), CompLevels);

       /* for ( const int& x : v[i] )
            cout << x << ' ';
        cout << '\n';   */
    }
   // cin.get();

    for ( int i = 1; i <= N; ++i )
        Update(i, a[i]);
}

void Update(int node, int val)
{
    trees[ind[node]].Update(FindPos(node), val);
}

int Query(int p1, int p2)
{
    if ( p1 == -1 || p2 == -1 )
        return 0;

    if ( ind[p1] == ind[p2] )
    {
        return trees[ind[p1]].Query( min(FindPos(p1), FindPos(p2)), max(FindPos(p1), FindPos(p2)) );
    }

    if ( lv[v[ind[p1]][0]] < lv[v[ind[p2]][0]] )
        swap(p1, p2);

    int v1 = trees[ind[p1]].Query( 0, FindPos(p1) );
    int v2 = Query(f[v[ind[p1]][0]], p2);

    return max(v1, v2);
}

void DFS( int node )
{
    if ( G[node].empty() )
    {
        treeSize.push_back(1);
        ++cntTrees;
        ind[node] = cntTrees - 1;
        cnt[node] = 1;

        return;
    }

    for ( const int& son : G[node] )
    {
        lv[son] = lv[node] + 1;
        DFS(son);
    }

    int maxim{0};
    cnt[node] = 1;
    for ( const int& son : G[node] )
    {
        cnt[node] += cnt[son];
        if ( cnt[son] > cnt[maxim] )
            maxim = son;
    }

    ind[node] = ind[maxim];
    treeSize[ind[node]]++;
}

Tree::Tree(int N1, int N2)
{
    int mid = (N1 + N2) / 2;
    i1 = N1, i2 = N2;
    val = 0;
    leftNode = nullptr;
    rightNode = nullptr;

    if ( N1 != N2 )
    {
        leftNode = new Tree(N1, mid);
        rightNode = new Tree(mid + 1, N2);
    }
}

int FindPos(int node)
{
    int st = 0, dr = v[ind[node]].size() - 1, mij;
    while ( st <= dr )
    {
        mij = ( st + dr ) / 2;

        if ( lv[v[ind[node]][mij]] < lv[node] )
            st = mij + 1;
        else
            if ( lv[v[ind[node]][mij]] > lv[node] )
                dr = mij - 1;
            else
                return mij;
    }

    return -1;
}

void Tree::Update(int pos, int value)
{
    if ( i1 == i2 )
    {
        val = value;

        return;
    }

    int mid = ( i1 + i2 ) / 2;
    if ( pos <= mid )
        leftNode->Update(pos, value);
    else
        rightNode->Update(pos, value);
    val = max(leftNode->val, rightNode->val);
}

int Tree::Query(int p1, int p2)
{
    if ( i1 >= p1 && i2 <= p2 )
    {
        return val;
    }

    int mid = (i1 + i2) / 2, t1{0}, t2{0};
    if ( p1 <= mid )
        t1 = leftNode->Query(p1, p2);
    if ( p2 > mid )
        t2 = rightNode->Query(p1, p2);
    return max(t1, t2);
}