Cod sursa(job #2498788)

Utilizator mitapirvuetAndrei Iachim mitapirvuet Data 24 noiembrie 2019 13:54:05
Problema Lowest Common Ancestor Scor 0
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 3.7 kb
#include <iostream>
#include <fstream>
#include <string.h>

#define N 100000
#define MAXLEVEL 100 // computed as ceil(log(N))
  
/**
 * DFS traversal to compute the depth for all the nodes and generate the 1st order parent for all nodes.
 * 
 * @currentNode: the current Node we are visiting
 * @prevNode: the previous Node we visited ( to store the parent )
 * @depth: vector that stores the depth for all nodes
 * @parent: vector that stores parents for level 2^i
 * @graph: The graph is represented using adjacency lists
 * 
 */
void dfs(int currentNode, int prevNode, std::vector<int>& depth,
         std::vector<std::vector<int>>& parent, const std::vector<std::vector<int>>& graph) {   
    parent[currentNode][0] = prevNode; 
    depth[currentNode] = depth[prevNode] + 1;

    for (int i = 0; i < graph[currentNode].size(); i++) { 
        if (graph[currentNode][i] != prevNode) {
            dfs(graph[currentNode][i], currentNode, depth, parent, graph); 
        } 
    } 
} 
  
/**
 * Compute 2^i parents for all nodes, with i in 0..MAXLEVEL.
 * 
 * @n: number of nodes
 * @parent: 2D array where element (i,j) represents parent on the 2^j level of node i.
 * 
 */
void getAllParents(int n, std::vector<std::vector<int>>& parent) {   
    for (int level = 1; level < MAXLEVEL; level++) { 
        for (int i = 1; i <= n; i++) { 
            if (parent[i][level-1] != -1) {
                parent[i][level] = parent[parent[i][level-1]][level-1]; 
            } 
        } 
    } 
} 

std::vector<int> lca(const std::vector<std::vector<int>>& graph, const std::vector< std::pair<int, int> >& queries) {
    std::vector<int> depth(graph.size(), 0);
    std::vector<std::vector<int>> parent(graph.size() + 1);

    for(auto& p:parent) {
        p.resize(MAXLEVEL, -1);
    } 

    // calculate depth for nodes
    dfs(1, 0, depth, parent, graph); 
  
    // calculate the node parents
    getAllParents(graph.size(), parent); 

    std::vector<int> results;

    for(auto query:queries) {
        int v = query.first, u = query.second;
        
        if (depth[v] < depth[u]) {
            std::swap(u, v); 
        }
    
        int h = depth[v] - depth[u]; 
    
        // get nodes to the same height, we know v is the deeper node
        for (int i = 0; i < MAXLEVEL; i++) {
            if (h % 2 == 1) {
                v = parent[v][i]; 
            }
            
            h /= 2;
        } 
    
        // edge case, if we reach the same node
        if (u == v) {
            results.push_back(u);
            continue;
        }
    
        // go level by level until LCA is found (we know this will take at most LOG N steps) 
        for (int i = MAXLEVEL -1 ; i >= 0; i--) {
            if (parent[u][i] != parent[v][i]) { 
                u = parent[u][i]; 
                v = parent[v][i]; 
            }
        } 
    
        results.push_back(parent[u][0]);
    }

    return results; 
} 

void addEdge(int a,int b, std::vector<std::vector<int>>& graph) 
{ 
    graph[a].push_back(b); 
    graph[b].push_back(a); 
} 

int main() 
{       
    std::ifstream inputStream;
    std::ofstream outputStream;

    int n, m;
    
    inputStream.open("lca.in");

    inputStream>>n>>m;

    std::vector<std::vector<int>> graph(n + 5);
    std::vector< std::pair<int, int> > queries;
    int src, dest;

    for(int i = 1; i < n ; ++i) {
        inputStream>>dest;

        addEdge(dest, i +1, graph);
    }

    int u , v;

    for(int i = 0; i < m; ++i) {
        inputStream>>u>>v;
        queries.push_back(std::make_pair(u, v));
    }

    std::vector<int> res = lca(graph, queries);

    outputStream.open("lca.out");

    for(auto& r:res) {
        outputStream<<r<<'\n';
    }

    outputStream.close();
    inputStream.close();
    

    return 0; 
}