CSES Solutions - Company Queries II

Last Updated : 14 Jun, 2024

A company has n employees, who form a tree hierarchy where each employee has a boss, except for the general director. Your task is to process q queries of the form: who is the lowest common boss of employees a and b in the hierarchy?

The employees are numbered 1,2,…,n, and employee 1 is the general director. The array arr[] has n−1 integers e2,e3,…,en: for each employee 2,3,…,n their boss. Finally, there are q queries, each query has two integers a and b: who is the lowest common boss of employees a and b?

Examples:

Input: n = 5, q = 3, arr[] = {1, 1, 3, 3}, queries[][] = {{4, 5}, {2, 5}, {1, 4}}
Output:
3
1
1
Explanation:

Lowest Common Boss of 4 and 5 is 3.
Lowest Common Boss of 2 and 5 is 1.
Lowest Common Boss of 1 and 4 is 1.
Input: n = 10, q = 3, arr[] = {1, 1, 1, 1, 2, 3, 4, 4, 1}, queries[][] = {{1, 8}, {2, 7}, {8, 3}}
Output:
1
1
1
Explanation:

Lowest Common Boss of 1 and 8 is 1.
Lowest Common Boss of 2 and 7 is 1.
Lowest Common Boss of 8 and 3 is 1.

Approach: To solve the problem, follow the below idea:

The problem can be solved using Binary Lifting Technique. We can construct the LCA table, such that LCA[i][j] stores the (2^j)th ancestor of node i. In order to fill the first column of LCA table, update the first column of each node i with immediate boss of node i. To fill the remaining columns, we can use the formula: LCA[i][j] = LCA[LCA[i][j - 1]][j - 1]. After filling the LCA[][] table completely, we can answer each query in logN time using Binary Lifting.

Step-by-step algorithm:

Run a Depth-First-Search to find the depth of all the nodes.
Fill the first column for every node i with the immediate boss of node i.
Fill the remaining columns with the formula: LCA[i][j] = LCA[LCA[i][j - 1]][j - 1].
To answer each query,
- Find the difference in depth of both the nodes.
- From the deeper node, use the LCA[][] table to make jumps in upward direction till both the nodes reach the same depth.
- After reaching the same depth, make jumps from both nodes till their ancestors are different.
- When the ancestor becomes same, return the ancestor node.

Below is the implementation of the algorithm:

C++

#include <bits/stdc++.h>
using namespace std;

const int MAXN = 200000;
const int LOG = 18; // because 2^18 > 200000

vector<int> adj[MAXN + 1]; // adjacency list of the tree
int LCA[MAXN + 1]
       [LOG]; // LCA[i][j] stores the 2^j-th ancestor of i
int depth[MAXN + 1]; // depth of each node

// DFS to find the depth of all the nodes and initialize the
// LCA table
void dfs(int v, int p)
{
    LCA[v][0] = p;
    // Fill the LCA table
    for (int i = 1; i < LOG; ++i) {
        if (LCA[v][i - 1] != -1) {
            LCA[v][i] = LCA[LCA[v][i - 1]][i - 1];
        }
        else {
            LCA[v][i] = -1;
        }
    }
    // Explore the neoghboring nodes
    for (int u : adj[v]) {
        if (u != p) {
            depth[u] = depth[v] + 1;
            dfs(u, v);
        }
    }
}

// function to find the LCA of node a and b
int lca(int a, int b)
{
    if (depth[a] < depth[b]) {
        swap(a, b);
    }

    int diff = depth[a] - depth[b];
    for (int i = 0; i < LOG; ++i) {
        if ((diff >> i) & 1) {
            a = LCA[a][i];
        }
    }

    if (a == b) {
        return a;
    }

    for (int i = LOG - 1; i >= 0; --i) {
        if (LCA[a][i] != LCA[b][i]) {
            a = LCA[a][i];
            b = LCA[b][i];
        }
    }

    return LCA[a][0];
}

int main()
{
    int n = 5, q = 3;
    int arr[] = { 1, 1, 3, 3 };
    vector<int> parent(n + 1);
    vector<vector<int> > queries
        = { { 4, 5 }, { 2, 5 }, { 1, 4 } };

    for (int i = 1; i < n; i++) {
        parent[i + 1] = arr[i - 1];
    }
    parent[1] = -1; // General director has no boss

    for (int i = 2; i <= n; ++i) {
        adj[parent[i]].push_back(i);
        adj[i].push_back(parent[i]);
    }

    depth[1] = 0;
    dfs(1, -1);

    for (int i = 0; i < q; ++i) {
        int a = queries[i][0], b = queries[i][1];
        cout << lca(a, b) << endl;
    }

    return 0;
}

Java

import java.util.*;

public class LCAFinder {

    static final int MAXN = 200000;
    static final int LOG = 18; // because 2^18 > 200000

    static List<Integer>[] adj = new ArrayList[MAXN + 1];
    static int[][] LCA = new int[MAXN + 1][LOG];
    static int[] depth = new int[MAXN + 1];

    // DFS to find the depth of all the nodes and initialize
    // the LCA table
    static void dfs(int v, int p)
    {
        LCA[v][0] = p;
        // Fill the LCA table
        for (int i = 1; i < LOG; ++i) {
            if (LCA[v][i - 1] != -1) {
                LCA[v][i] = LCA[LCA[v][i - 1]][i - 1];
            }
            else {
                LCA[v][i] = -1;
            }
        }
        // Explore the neighboring nodes
        for (int u : adj[v]) {
            if (u != p) {
                depth[u] = depth[v] + 1;
                dfs(u, v);
            }
        }
    }

    // function to find the LCA of node a and b
    static int lca(int a, int b)
    {
        if (depth[a] < depth[b]) {
            int temp = a;
            a = b;
            b = temp;
        }

        int diff = depth[a] - depth[b];
        for (int i = 0; i < LOG; ++i) {
            if ((diff >> i & 1) == 1) {
                a = LCA[a][i];
            }
        }

        if (a == b) {
            return a;
        }

        for (int i = LOG - 1; i >= 0; --i) {
            if (LCA[a][i] != LCA[b][i]) {
                a = LCA[a][i];
                b = LCA[b][i];
            }
        }

        return LCA[a][0];
    }

    public static void main(String[] args)
    {
        int n = 5, q = 3;
        int[] arr = { 1, 1, 3, 3 };
        int[] parent = new int[n + 1];
        List<int[]> queries = Arrays.asList(
            new int[] { 4, 5 }, new int[] { 2, 5 },
            new int[] { 1, 4 });

        for (int i = 0; i <= MAXN; ++i) {
            adj[i] = new ArrayList<>();
        }

        for (int i = 1; i < n; i++) {
            parent[i + 1] = arr[i - 1];
        }
        parent[1] = -1; // General director has no boss

        for (int i = 2; i <= n; ++i) {
            adj[parent[i]].add(i);
            adj[i].add(parent[i]);
        }

        depth[1] = 0;
        dfs(1, -1);

        for (int[] query : queries) {
            int a = query[0], b = query[1];
            System.out.println(lca(a, b));
        }
    }
}

// This code is contributed by Shivam Gupta

Output

3
1
1

Time Complexity: O(nlogn + qlogn)
Auxiliary Space: O(nlogn)

Cognizant SDE Sheet: Interview Questions and Answers

madhavwvcvq

Improve

Article Tags :

Practice Tags :

Graph

CSES Solutions - Company Queries II

Similar Reads

Thank You!

What kind of Experience do you want to share?