Cod sursa(job #2982608)

Utilizator ArseniuVictorStefanArseniu Victor Stefan ArseniuVictorStefan Data 20 februarie 2023 15:53:31
Problema Heavy Path Decomposition Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 6.64 kb
#include <stdio.h>

// no edge
#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
const int& max(const int& x, const int& y)
{
    return x >= y ? x : y;
}

// segment tree for the HPD
struct
{
    // the tree
    int tree[2 * p2(NMAX)];
    int sz;

    // init the tree with n leafs
    void init(int n)
    {
        sz = p2(n);
    }

    // set the i-th value to x
    void set(int i, int x)
    {
        tree[i + sz] = x;
    }

    // update the tree with x on position i
    void update(int i, int x)
    {
        i += sz;
        tree[i] = x;

        while(i >>= 1)
            tree[i] = max(tree[i << 1], tree[i << 1 | 1]);
    }

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

    // query 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;
    }
} st;


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

    // the nodes
    struct
    {
        int value; // the value of the node
        int begin; // the first edge in the linked list of edges
        int parent; // the parent of the node
        int depth; // the depth of the node
        int size; // the size of the subtree of the node
        int heavy; // the heavy child
        int head; // the head of the heavy path the node is on
        int index; // the index of the node in the segment tree
    } nodes[NMAX];

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

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

    // adds 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 "head"
    void decompose(int head)
    {
        int u = head;
        while(u != NONE)
        {
            nodes[u].head = head;
            nodes[u].index = nr;
            st.set(nr++, nodes[u].value);
            
            u = nodes[u].heavy;
        }
    }

    // dfs for heavy path decomposition
    void dfs(int u)
    {
        nodes[u].size = 1;

        for(int i = nodes[u].begin; i != NONE; i = edges[i].next)
        {
            int v = edges[i].v;
            if(v != nodes[u].parent)
            {
                nodes[v].parent = u;
                nodes[v].depth = nodes[u].depth + 1;
                
                dfs(v);

                nodes[u].size += nodes[v].size;

                if(nodes[u].heavy == NONE)
                    nodes[u].heavy = v;
                else if(nodes[nodes[u].heavy].size < nodes[v].size)
                {
                    decompose(nodes[u].heavy);
                    nodes[u].heavy = v;
                }
                else decompose(v);
            }
        }
    }

    // heavy path decomposition with u as the root
    void hpd(int u)
    {
        nodes[u].depth = 0;
        dfs(u);
        decompose(u);
    }

    // find the max. value on the path from x to y
    int query(int x, int y)
    {
        int ans = 0;
        while(nodes[x].head != nodes[y].head)
        {
            // prefix query on the lower path
            if(nodes[nodes[x].head].depth < nodes[nodes[y].head].depth)
            {
                int tmp = x;
                x = y;
                y = tmp;
            }

            int q = st.query(nodes[nodes[x].head].index, nodes[x].index + 1);
            //fprintf(stderr, "Query: %d %d -> %d\n", nodes[x].head, x, q);

            ans = max(ans, q);

            // move up the path
            x = nodes[nodes[x].head].parent;
        }

        // query on the last path
        if(nodes[x].index > nodes[y].index)
        {
            int tmp = x;
            x = y;
            y = tmp;
        }

        int q = st.query(nodes[x].index, nodes[y].index + 1);
        //fprintf(stderr, "Query: %d %d -> %d\n", x, y, q);

        ans = max(ans, q);

        // return the answer
        return ans;
    }
    
} graph;


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

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

    // read the graph and no. of queries
    scanf("%d %d", &n, &m);

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

    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
    st.init(n);

    // heavy path decomposition from root 0
    graph.hpd(0);

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

    // spit out the whole computation
    /*
    fprintf(stderr, "Index:");
    for(int i = 0; i < n; i++)
        fprintf(stderr, " %d", i);
    fprintf(stderr, "\n");
    fprintf(stderr, "Heavy:");
    for(int i = 0; i < n; i++)
    {
        if(graph.nodes[i].heavy != NONE)
            fprintf(stderr, " %d", graph.nodes[i].heavy);
        else fprintf(stderr, " N");
    }
    fprintf(stderr, "\n");
    fprintf(stderr, "Head: ");
    for(int i = 0; i < n; i++)
    {
        if(graph.nodes[i].head != NONE)
            fprintf(stderr, " %d", graph.nodes[i].head);
        else fprintf(stderr, " N");
    }
    fprintf(stderr, "\n");
    fprintf(stderr, "Index:");
    for(int i = 0; i < n; i++)
    {
        if(graph.nodes[i].index != NONE)
            fprintf(stderr, " %d", graph.nodes[i].index);
        else fprintf(stderr, " N");
    }
    fprintf(stderr, "\n");
    */

    // answer the queries
    for(int i = 0; i < m; i++)
    {
        int t, x, y;
        scanf("%d %d %d", &t, &x, &y);
        if(t == 0)
        {
            x--;
            st.update(graph.nodes[x].index, y);
        }
        else
        {
            x--; y--;
            printf("%d\n", graph.query(x, y));
        }
    }

    // exit
    return 0;
}