Cod sursa(job #2982484)

Utilizator ArseniuVictorStefanArseniu Victor Stefan ArseniuVictorStefan Data 20 februarie 2023 12:52:39
Problema Heavy Path Decomposition Scor 0
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 8.81 kb
#include <stdio.h>

// no node
#define NONE (-1)

// the max. no. of nodes
#define NMAX (100'000)

// returns the smallest power of 2 bigger than or equal to n
constexpr int p2(int n)
{
    int p = 1;
    while(p < n)
        p <<= 1;
    return p;
}

// returns the max. of x and y
constexpr const int& max(const int& x, const int& y)
{
    return x >= y ? x : y;
}

// segment tree for hpd
struct
{
    // the tree itself and the size of the tree
    int tree[2 * p2(NMAX)];
    int sz;

    // create a tree with at least n nodes
    void create(int n)
    {
        // set the size
        sz = p2(n);
    }

    // set the node at i to x without updating the tree
    void set(int i, int x)
    {
        tree[sz + i] = x;
    }

    // build the segment tree
    void build()
    {
        for(int i = sz - 1; i >= 1; i--)
            tree[i] = max(tree[i << 1], tree[i << 1 | 1]);
    }

    // update the segment tree at i to x
    void update(int i, int x)
    {
        // set the first value
        i += sz;
        tree[i] = x;

        // update the upper layers
        while(i >>= 1)
            tree[i] = max(tree[i << 1], tree[i << 1 | 1]);
    }

    // query the answer on the interval [l, r)
    int query(int l, int r)
    {
        int ans = 0;
        for(l += sz, r += sz; l < r; l >>= 1, r >>= 1)
        {
            if(l & 1)
                ans = max(ans, tree[l++]);
            if(r & 1)
                ans = max(ans, tree[--r]);
        }
        return ans;
    }
} segment_tree;

// the graph
struct
{
    // the edges of the graph
    struct
    {
        int v;
        int next;
    } edges[2 * (NMAX - 1)];

    // the nodes of the graph
    struct
    {
        int value; // the value of the node
        int begin; // the first edge
        int size; // the size of its subtree
        int heavy; // the heavy child of the node
        int path_parent; // the parent of the path the node is on
        int path; // the beginning of the heavy path the node is on
        int index; // the index on the heavy path

        int parent; // the parent of the node
        int depth; // the depth of the node
        int lifting; // the lifting edge

        int start; // the DFS start time of the node
        int end; // the DFS end time of the node
    } nodes[NMAX];

    // the no. of nodes and no. of edges
    int n, m;

    // add a node with value v
    void add_node(int v)
    {
        nodes[n++] = { v, NONE, NONE, NONE, NONE, NONE, NONE, NONE, NONE, NONE, NONE, NONE };
    }

    // add an edge from u to v
    void add_edge(int u, int v)
    {
        edges[m] = { v, nodes[u].begin };
        nodes[u].begin = m++;
    }

    // the no. of nodes already decomposed
    int nr;

    // decompose the heavy path starting at u with the given parent
    void decompose(int u, int parent)
    {
        // init the path and the index
        nodes[u].path_parent = parent;
        nodes[u].path = nr;
        nodes[u].index = nr;
        segment_tree.set(nr, nodes[u].value);
        nr++;

        // go down the path and set the indices
        while(nodes[u].heavy != NONE)
        {
            int v = nodes[u].heavy;

            // set the path and the index of v
            nodes[v].path_parent = parent;
            nodes[v].path = nodes[u].path;
            nodes[v].index = nr;
            segment_tree.set(nr, nodes[v].value);
            nr++;

            // move down the path
            u = v;
        }
    }

    // the current DFS time
    int time;

    // decompose the subtree of u (all the paths that are complete in this subtree)
    // also, calculate the binary lifting paths for LCA computation
    // also, calculate start and end times of DFS for ancestor check
    void heavy(int u)
    {
        // if the parent isn't NONE, calculate the lifting
        if(nodes[u].parent != NONE)
        {
            int v = nodes[u].parent;
            int w = nodes[v].lifting;
            int x = nodes[w].lifting;

            // if the distances are equal, set the parent
            if(nodes[x].depth - nodes[w].depth == nodes[w].depth - nodes[v].depth)
                nodes[u].lifting = x;
            else nodes[u].lifting = v;
        }
        // else, the node is the root, so init depth and lifting accordingly
        else
        {
            nodes[u].depth = 0;
            nodes[u].parent = u;
            nodes[u].lifting = u;
        }

        // set the start time
        nodes[u].start = time++;

        // the heavy child of u
        int h = NONE;

        // init the size of u
        nodes[u].size = 1;

        // go through the neighbours and recursive call
        for(int i = nodes[u].begin; i != NONE; i = edges[i].next)
        {
            int v = edges[i].v;

            // if v wasn't visited i.e. its depth is still NONE
            if(nodes[v].parent == NONE)
            {
                // set the depth, parent; recursive call and update the size of u
                nodes[v].depth = nodes[u].depth + 1;
                nodes[v].parent = u;
                heavy(v);
                nodes[u].size += nodes[v].size;

                // update the heavy edge
                if(h == NONE)
                    h = v;
                else if(nodes[v].size > nodes[h].size)
                {
                    // decompose the path starting at h becaure h isn't heavy anymore so a path must start at it
                    decompose(h, u);
                    h = v;
                }
                // else, v is light so a heavy path must start at it, so decompose that path
                else decompose(v, u);
            }
        }

        // set the heavy edge of u
        nodes[u].heavy = h;

        // set the end time
        nodes[u].end = time;
    }

    // returns true if u is an ancestor of v
    bool is_ancestor(int u, int v)
    {
        return nodes[u].start <= nodes[v].start && nodes[v].start <= nodes[u].end;
    }

    // calculate the LCA of u and v
    int lca(int u, int v)
    {
        // make sure start[u] <= start[v]
        if(nodes[u].start > nodes[v].start)
        {
            int tmp = u;
            u = v;
            v = tmp;
        }

        // while u isn't an ancestor of v, raise it
        while(!is_ancestor(u, v))
        {
            // if we can use the lifting, use it
            if(!is_ancestor(nodes[u].lifting, v))
                u = nodes[u].lifting;
            // else, just go to the parent
            else u = nodes[u].parent;
        }

        // return the LCA
        return u;
    }

    // query the path between x and y where y is guaranteed to be an ancestor of x
    int lca_query(int x, int y)
    {
        // the answer
        int ans = 0;

        // while the nodes aren't on the same path, prefix queries and move up the path
        while(nodes[x].path != nodes[y].path)
        {
            ans = max(ans, segment_tree.query(nodes[x].path, nodes[x].index + 1));
            x = nodes[x].path_parent;
        }

        // take into account the nodes left between the indices of x and y on their common path
        ans = max(ans, segment_tree.query(nodes[y].index, nodes[x].index + 1));

        // return the answer
        return ans;
    }

    // query the path between x and y
    int query(int x, int y)
    {
        int l = lca(x, y);
        return max(lca_query(x, l), lca_query(y, l));
    }
} graph;

// the no. of nodes and no. of operations
int n, m;

int main()
{
    // open the files
    freopen("heavypath.in", "r", stdin);
    freopen("heavypath.out", "w", stdout);

    // read the input
    scanf("%d %d", &n, &m);

    // read the nodes
    for(int i = 0; i < n; i++)
    {
        int v;
        scanf("%d", &v);
        graph.add_node(v);
    }

    // read the edges
    for(int i = 1; i < n; i++)
    {
        int u, v;
        scanf("%d %d", &u, &v);
        u--; v--;

        graph.add_edge(u, v);
        graph.add_edge(v, u);
    }

    // init the segment tree
    segment_tree.create(n);

    // build the heavy path decomposition
    graph.heavy(0);

    // decompose the last heavy path
    graph.decompose(0, NONE);

    // build the segment tree
    segment_tree.build();

    // read the queries and answer them
    for(int i = 0; i < m; i++)
    {
        int t, x, y;
        scanf("%d %d %d", &t, &x, &y);

        // update
        if(t == 0)
        {
            x--;
            segment_tree.update(graph.nodes[x].index, y);
        }
        // query
        else
        {
            x--; y--;

            // print the query
            printf("%d\n", graph.query(x, y));
        }
    }

    // exit
    return 0;
}