Cod sursa(job #3170272)

Utilizator AleXutzZuDavid Alex Robert AleXutzZu Data 17 noiembrie 2023 09:40:15
Problema Heavy Path Decomposition Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 3.81 kb
#include <iostream>
#include <fstream>
#include <vector>

struct SegmentTree {
private:
    std::vector<int> tree;

    int query(int left, int right, int node, int start, int end) {
        if (start <= left && right <= end) return tree[node];

        if (end < left || right < start) return 0;
        int mid = (left + right) / 2;
        int ql = query(left, mid, 2 * node, start, end);
        int qr = query(mid + 1, right, 2 * node + 1, start, end);
        return std::max(ql, qr);
    }

    void update(int left, int right, int node, int pos, int value) {
        if (left == right) tree[node] = value;
        else {
            int mid = (left + right) / 2;
            if (pos <= mid) update(left, mid, 2 * node, pos, value);
            else update(mid + 1, right, 2 * node + 1, pos, value);
            tree[node] = std::max(tree[2 * node], tree[2 * node + 1]);
        }
    }

    int size;
public:
    explicit SegmentTree(int size) : size(size) {
        tree.resize(4 * size + 1);
    }

    void update(int pos, int value) {
        update(1, size, 1, pos, value);
    }

    int query(int start, int end) {
        return query(1, size, 1, start, end);
    }
};

struct Graph {
private:
    std::vector<std::vector<int>> adj;
    std::vector<int> value;

    std::vector<int> parent;
    std::vector<int> depth;
    std::vector<int> heavy;
    std::vector<int> pos;
    std::vector<int> subtree;
    int ptr = 0;
    int size;
    SegmentTree tree;

    void calc(int node, int root = 1) {
        depth[node] = 1 + depth[root];
        parent[node] = root;
        subtree[node] = 1;

        for (const auto &x: adj[node]) {
            if (x == root) continue;
            calc(x, node);
            subtree[node] += subtree[x];
        }
    }

    void decompose(int node, int h_node, int root = 1) {
        pos[node] = ++ptr;
        heavy[node] = h_node;

        tree.update(pos[node], value[node]);

        int h_child = -1, max_size = 0;
        for (const auto &x: adj[node]) {
            if (x == root) continue;
            if (subtree[x] > max_size) max_size = subtree[x], h_child = x;
        }

        if (h_child == -1) return;
        decompose(h_child, h_node, node);

        for (const auto &x: adj[node]) {
            if (x == root || x == h_child) continue;
            decompose(x, x, node);
        }
    }

public:
    explicit Graph(std::vector<int> &values) : size((int) values.size() - 1), value(values), tree(size) {
        adj.resize(size + 1);
        parent.resize(size + 1);
        depth.resize(size + 1);
        heavy.resize(size + 1);
        pos.resize(size + 1);
        subtree.resize(size + 1);
    }

    void add_edge(int x, int y) {
        adj[x].push_back(y);
        adj[y].push_back(x);
    }

    void heavy_light() {
        calc(1);
        decompose(1, 1);
    }

    void update(int node, int val) {
        tree.update(pos[node], val);
    }

    int query(int x, int y) {
        int ans = 0;
        while (heavy[x] != heavy[y]) {
            if (depth[heavy[x]] < depth[heavy[y]]) std::swap(x, y);
            ans = std::max(ans, tree.query(pos[heavy[x]], pos[x]));
            x = parent[heavy[x]];
        }

        if (depth[x] > depth[y]) std::swap(x, y);
        ans = std::max(ans, tree.query(pos[x], pos[y]));
        return ans;
    }
};

int main() {
    std::ifstream input("heavypath.in");
    std::ofstream output("heavypath.out");

    int n, q;
    input >> n >> q;
    std::vector<int> values(n + 1);
    for (int i = 1; i <= n; ++i) input >> values[i];
    Graph tree(values);

    for (int i = 0; i < n - 1; ++i) {
        int x, y;
        input >> x >> y;
        tree.add_edge(x, y);
    }

    tree.heavy_light();

    while (q--) {
        int t, x, y;
        input >> t >> x >> y;
        if (t == 0) tree.update(x, y);
        else output << tree.query(x, y) << '\n';
    }
    return 0;
}