Cod sursa(job #2806300)

Utilizator marceldosmionsMarcel marceldosmions Data 22 noiembrie 2021 14:59:58
Problema Lowest Common Ancestor Scor 60
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 3.09 kb
#include <stack>
#include <vector>
#include <fstream>

using namespace std;

class disjoint_set {
  int size;
  vector<int> parent, rank;

public:
  disjoint_set(int size) {
    this->size = size;
    parent.resize(size, -1);
    rank.resize(size);
  }

  void make_set(int x) {
    parent[x] = x;
    rank[x] = 1;
  }

  int find(int x) {
    if(parent[x] != x){
      parent[x] = find(parent[x]);
    }

    return parent[x];
  }

  void Union(int x, int y) {
    int xset = find(x), yset = find(y);
    if (xset == yset) {
      return;
    }

    if (rank[xset] > rank[yset]) {
      parent[yset] = xset;
    } else if (rank[xset] < rank[yset]){
      parent[xset] = yset;
    }

    if (rank[xset] == rank[yset]) {
      parent[xset] = yset;
      rank[yset]++;
    }
  }

  void set_ancestors(int x, int ancestor) {
    int root = find(x);
    parent[root] = ancestor;
    parent[ancestor] = ancestor;
    if (rank[root] > rank[ancestor]) {
      rank[ancestor] = rank[root] + 1;
    }
  }
};

static pair<int, int> order_pair(pair<int, int> p) {
  return {min(p.first, p.second), max(p.first, p.second)};
}

vector<int> lca(const vector<vector<int>> &tree, const vector<pair<int, int>> &queries) {
  vector<vector<int>> qst(tree.size());
  vector<int> res(queries.size()), app(tree.size(), 0);
  stack<pair<int, int>> st;
  pair<int, int> query_pair;
  disjoint_set forest(tree.size());
  for (int i = 0; i < tree.size(); i++) {
    forest.make_set(i);
  }
  for (int i = 0; i < queries.size(); i++) {
    qst[queries[i].first].push_back(i);
    qst[queries[i].second].push_back(i);
  }
  st.push(make_pair(0, 1));
  while (!st.empty()) {
    int node = st.top().first;
    int index = st.top().second;
    if (index == tree[node].size()) {
      app[node] = 1;
      for (const int &pos : qst[node]) {
        query_pair = order_pair(queries[pos]);
        int other = (node == query_pair.first) ? query_pair.second : query_pair.first;
        if (app[other]) {
          res[pos] = forest.find(other);
        }
      }
      forest.Union(node, tree[node][0]);
      forest.set_ancestors(node, tree[node][0]);
      st.pop();
    } else {
      st.top().second++;
      st.push(make_pair(tree[node][index], 1));
    }
  }

  return res;
}

void read_data(vector<vector<int>> &tree, vector<pair<int, int>> &queries){
    ifstream input;
    input.open("./lca.in", ios::in);
    if(!input.is_open()) return;

    int n, m, node, n1, n2;
    input>>n>>m;
    tree.resize(n, vector<int>(1));

    for(int i = 2; i <= n; i++){
        input>>node;
        tree[i - 1][0] = node - 1;
        tree[node - 1].push_back(i - 1);
    }

    for(int i = 0; i < m; i++){
        input>>n1>>n2;
        queries.push_back({n1 - 1, n2 - 1});
    }
    
    input.close();
}

void print_data(vector<int> res){
    ofstream output;
    output.open("./lca.out", ios::out);
    if(!output.is_open()) return;

    for(int i = 0; i < res.size(); i++){
        output<<res[i] + 1<<endl;
    }

    output.close();
}

int main(){
    vector<vector<int>> tree;
    vector<pair<int, int>> queries;

    read_data(tree, queries);
    print_data(lca(tree, queries));

    return 0; 
}