Cod sursa(job #2986372)

Utilizator MatteoalexandruMatteo Verzotti Matteoalexandru Data 28 februarie 2023 13:38:38
Problema Heavy Path Decomposition Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 4.26 kb
#include <bits/stdc++.h>
#ifdef BLAT
    #include "debug/debug.hpp"
#else
    #define debug(x...)
#endif

using namespace std;

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

using Graph = vector<vector<int>>;

struct Node {
    int value;              // value of the node
    int parent, depth;      // parent and depth of each node
    int heavy, head, pos;   // index of heavy node, head node and position in the decomposition
};

namespace HeavyLight {
    // Calculates the size of each subtree and computes heavy nodes
    int dfs(int node, int parent, vector<Node> &nodes, const Graph &g) {
        int size = 1, maxChild = 0;
        nodes[node].parent = parent;
        nodes[node].depth = 1 + nodes[parent].depth;

        for (auto to : g[node])
            if (to != parent) {
                // Find the child subtree size
                int childSize = dfs(to, node, nodes, g);

                // If it's the largest child, update heavy node
                if (childSize > maxChild) {
                    maxChild = childSize;
                    nodes[node].heavy = to;
                }

                size += childSize;
            }

        return size;
    }

    int currpos = 0;
    vector<int> decomp = {0};
    // Performs the decomposition
    void decompose(int node, int head, vector<Node> &nodes, const Graph &g) {
        Node &currnode = nodes[node];
        currnode.head = head; currnode.pos = ++currpos;
        decomp.push_back(node);

        // If we have a heavy child we update that first
        if (currnode.heavy != 0)
            decompose(currnode.heavy, head, nodes, g);

        // We update the rest of the children
        for (auto to : g[node]) {
            if (to != currnode.parent && to != currnode.heavy) {
                // Start a new chain with the child as head node
                decompose(to, to, nodes, g);
            }
        }
    }
}

class SegTree {
private:
    int n;
    vector<int> a;

    void recalculate(int node) {
        a[node] = max(a[node << 1], a[node << 1 | 1]);
    }

    void _update(int node, int l, int r, int pos, int val) {
        if (l == r) {
            a[node] = val;
            return;
        }
        int mid = (l + r) >> 1;
        if (pos <= mid) _update(node << 1, l, mid, pos, val);
        else _update(node << 1 | 1, mid + 1, r, pos, val);

        recalculate(node);
    }

    int _query(int node, int l, int r, int x, int y) {
        if (x <= l && r <= y) 
            return a[node];
        int mid = (l + r) >> 1, ans = 0;
        if (x <= mid) ans = max(ans, _query(node << 1, l, mid, x, y));
        if (mid < y) ans = max(ans, _query(node << 1 | 1, mid + 1, r, x, y));

        return ans;
    }

public:
    SegTree(int _n) : n(_n) {
        a.resize(1 + 4 * n);
    }

    void update(int pos, int val) {
        _update(1, 1, n, pos, val);
    }

    int query(int x, int y) {
        return _query(1, 1, n, x, y);
    }
};

int query(int x, int y, const vector<Node> &nodes, SegTree &aint) {
    int ans = 0;

    // Try to get x and y on the same chain
    while (nodes[x].head != nodes[y].head) {
        if (nodes[nodes[x].head].depth > nodes[nodes[y].head].depth)
            swap(x, y);
        
        // Query on the chain
        ans = max(ans, aint.query(nodes[nodes[y].head].pos, nodes[y].pos));
        y = nodes[nodes[y].head].parent;
    }

    if (nodes[x].depth > nodes[y].depth) swap(x, y);
    ans = max(ans, aint.query(nodes[x].pos, nodes[y].pos));

    return ans;
}

int main() {
    int n, m;
    in >> n >> m;

    vector<Node> nodes(1 + n);
    Graph g(1 + n);
    for (int i = 1; i <= n; i++)
        in >> nodes[i].value;

    for (int i = 1; i < n; i++) {
        int x, y;
        in >> x >> y;
        g[x].push_back(y);
        g[y].push_back(x);
    }

    // Start the HLD
    HeavyLight::dfs(1, 0, nodes, g);
    HeavyLight::decompose(1, 1, nodes, g);
    debug(HeavyLight::decomp);

    // Build the segment tree
    SegTree aint(n);
    for (int i = 1; i <= n; i++)
        aint.update(nodes[i].pos, nodes[i].value);

    while (m--) {
        int t, x, y;
        in >> t >> x >> y;
        debug(t, x, y);
        if (t == 0)
            aint.update(nodes[x].pos, y);
        else out << query(x, y, nodes, aint) << '\n';
    }
    in.close();
    out.close();
    return 0;
}