LC 3108: Augmented Union-Find.

March 21, 2025

Find the problem description here. If you would like to read the solution in leetcode platform, find it here.


Intuition

The property of BITWISE AND operation is the key here.

When we apply BITWISE AND operation on a bunch of integers, then the answer can be equal to the minimum integer out of all or even lesser than the minimum integer. This is the only key to solve this problem. The problem asked us to minimize the cost. Therefore, the minimum comes from the all the integers of a connected component.

Everything else is same as the other Disjoint solutions following the intuition that we need to find all the connected components. And for any two vertices that are in the same component, the cost is same which is the BITWISE AND of all the edge weights in that component.

The only thing different here is that, I augmented the cost calculation into the union operation without calculating it afterwards, which is pretty simple.

For remaining everything, you can direcly look at the code.


Complexity

  • Time complexity: O(v+e+q)O(v + e + q), where v is the number of vertices, e is the number of edges and q is the number of queries.

  • Union-Find operation follows the Inverse Ackermann Function which is practically, can be taken as constant time with all the pruning find steps and union by rank efficiency.

  • Space complexity: O(n)O(n) for parent, depth and cost arrays.


Code

class Solution {
public:
    int find(vector<int>& parent, int a) {
        if (parent[a] == a) return a;
        return parent[a] = find(parent, parent[a]); // pruning for the later find operations
    }

    void union_(vector<int>& parent, vector<int>& depth, vector<uint>& cost, int a, int b, int weight) {
        int parentA = find(parent, a);
        int parentB = find(parent, b);
        
        // an optional step - union by rank for efficiency
        // if depth of parentA is greater than parentB, swap it with parentB
        // as we always assigning the parentB as the final parent
        if (depth[parentA] > depth[parentB])    swap(parentA, parentB);

        cost[parentB] &= weight;            // AND with current weight
        cost[parentB] &= cost[parentA];     // AND with the parentA cost
        parent[parentA] = parentB;          // assigning parentB as representative of both the components

        // if the depths are same, increment the depth of parentB
        if (depth[parentA] == depth[parentB])   depth[parentB]++;
    }

    vector<int> minimumCost(int n, vector<vector<int>>& edges, vector<vector<int>>& query) {
        vector<int> parent(n);
        vector<int> depth(n, 0);

        // you can take INT_MAX as well here which works for the constraints 10^5
        // or ~0 or (1 << 17) - 1 (a close value to 10^5)
        vector<uint> cost(n, -1);   // initial cost should be all set bits

        for (int i = 0; i < n; ++i) {
            parent[i] = i;  // initially everything is parent to itself
        }

        for (vector<int> edge : edges) {
            union_(parent, depth, cost, edge[0], edge[1], edge[2]);
        }

        vector<int> result;
        for (vector<int> q : query) {
            // if the parent of both the nodes are different, then they are not connected
            if (find(parent, q[0]) != find(parent, q[1]))   result.push_back(-1);
            // if they are connected, then we take the cost of the current component
            else    result.push_back(cost[find(parent, q[0])]);
        }

        return result;
    }
};

good luck!