Cod sursa(job #2501592)

Utilizator preda.andreiPreda Andrei preda.andrei Data 29 noiembrie 2019 22:56:24
Problema Lowest Common Ancestor Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 3.1 kb
#include <fstream>
#include <vector>

class DisjointSet {
public:
    explicit DisjointSet(int size) : father_(size, -1), marked_(size, false) {}

    int getRoot(int node) {
        if (father_[node] == -1) {
            return node;
        }
        return father_[node] = getRoot(father_[node]);
    }

    void unite(int a, int b) {
        int root_a = getRoot(a);
        int root_b = getRoot(b);

        if (root_a != root_b) {
            father_[root_b] = root_a;
        }
    }

    void mark(int node) { marked_[node] = true; }

    bool isMarked(int node) const { return marked_[node]; }

private:
    std::vector<int> father_;
    std::vector<bool> marked_;
};

struct Query {
    int other_node;
    int index;
    int answer;

    Query(int other_node, int index)
        : other_node(other_node), index(index), answer(-1) {}
};

using Tree = std::vector<std::vector<int>>;
using Queries = std::vector<std::vector<Query>>;

static Queries buildQueriesByNode(int nodes,
                           const std::vector<std::pair<int, int>> &queries) {
    Queries queries_by_node(nodes);
    for (size_t i = 0; i < queries.size(); i += 1) {
        int a = queries[i].first;
        int b = queries[i].second;

        queries_by_node[a].push_back(Query(b, i));
        queries_by_node[b].push_back(Query(a, i));
    }
    return queries_by_node;
}

static std::vector<int> extractAnswers(int count, const Queries &queries) {
    std::vector<int> answers(count);
    for (const auto &q_list : queries) {
        for (const auto &q : q_list) {
            if (q.answer != -1) {
                answers[q.index] = q.answer;
            }
        }
    }
    return answers;
}

static void dfs(const Tree &tree, int node, int father,
                DisjointSet &forest, Queries &queries) {
    for (const auto &son : tree[node]) {
        if (son != father) {
            dfs(tree, son, node, forest, queries);
            forest.unite(node, son);
        }
    }

    forest.mark(node);

    for (auto &q : queries[node]) {
        if (forest.isMarked(q.other_node)) {
            q.answer = forest.getRoot(q.other_node);
        }
    }
}

std::vector<int> lca(const Tree &tree,
                     const std::vector<std::pair<int, int>> &queries) {
    int nodes = tree.size();
    auto queries_by_node = buildQueriesByNode(nodes, queries);

    DisjointSet forest(nodes);
    dfs(tree, 0, -1, forest, queries_by_node);

    auto answers = extractAnswers(queries.size(), queries_by_node);
    return answers;
}

int main()
{
    std::ifstream fin("lca.in");
    std::ofstream fout("lca.out");

    int nodes, tests;
    fin >> nodes >> tests;

    std::vector<std::vector<int>> tree(nodes);
    for (int i = 1; i < nodes; i += 1) {
        int father;
        fin >> father;
        tree[father - 1].push_back(i);
        tree[i].push_back(father - 1);
    }

    std::vector<std::pair<int, int>> queries(tests);
    for (auto &q : queries) {
        fin >> q.first >> q.second;
        q.first -= 1;
        q.second -= 1;
    }

    auto res = lca(tree, queries);
    for (const auto &node : res) {
        fout << node + 1 << "\n";
    }

    return 0;
}