LeetCode-in-Java

3607. Power Grid Maintenance

Medium

You are given an integer c representing c power stations, each with a unique identifier id from 1 to c (1‑based indexing).

These stations are interconnected via n bidirectional cables, represented by a 2D array connections, where each element connections[i] = [ui, vi] indicates a connection between station ui and station vi. Stations that are directly or indirectly connected form a power grid.

Initially, all stations are online (operational).

You are also given a 2D array queries, where each query is one of the following two types:

Return an array of integers representing the results of each query of type [1, x] in the order they appear.

Note: The power grid preserves its structure; an offline (non‑operational) node remains part of its grid and taking it offline does not alter connectivity.

Example 1:

Input: c = 5, connections = [[1,2],[2,3],[3,4],[4,5]], queries = [[1,3],[2,1],[1,1],[2,2],[1,2]]

Output: [3,2,3]

Explanation:

Example 2:

Input: c = 3, connections = [], queries = [[1,1],[2,1],[1,1]]

Output: [1,-1]

Explanation:

Constraints:

Solution

import java.util.ArrayList;
import java.util.List;
import java.util.PriorityQueue;

@SuppressWarnings("unchecked")
public class Solution {
    private static class UF {
        int[] par;
        PriorityQueue<Integer>[] pq;
        boolean[] active;

        UF(int n) {
            par = new int[n];
            pq = new PriorityQueue[n];
            active = new boolean[n];
            for (int i = 0; i < n; i++) {
                active[i] = true;
                par[i] = i;
                pq[i] = new PriorityQueue<>();
                pq[i].add(i);
            }
        }

        int find(int u) {
            if (par[u] == u) {
                return u;
            }
            par[u] = find(par[u]);
            return par[u];
        }

        void union(int u, int v) {
            int pu = find(u);
            int pv = find(v);
            if (pu == pv) {
                return;
            }
            if (pq[pu].size() > pq[pv].size()) {
                while (!pq[pv].isEmpty()) {
                    pq[pu].add(pq[pv].poll());
                }
                par[pv] = par[pu];
            } else {
                while (!pq[pu].isEmpty()) {
                    pq[pv].add(pq[pu].poll());
                }
                par[pu] = par[pv];
            }
        }

        void inactive(int u) {
            active[u] = false;
        }

        int check(int u) {
            if (active[u]) {
                return u;
            }
            int pu = find(u);
            while (!pq[pu].isEmpty() && !active[pq[pu].peek()]) {
                pq[pu].poll();
            }
            return !pq[pu].isEmpty() ? pq[pu].peek() : -2;
        }
    }

    public int[] processQueries(int c, int[][] connections, int[][] queries) {
        UF uf = new UF(c);
        for (int[] con : connections) {
            int u = con[0];
            int v = con[1];
            uf.union(u - 1, v - 1);
        }
        List<Integer> res = new ArrayList<>();
        for (int[] q : queries) {
            if (q[0] == 1) {
                res.add(uf.check(q[1] - 1) + 1);
            } else {
                uf.inactive(q[1] - 1);
            }
        }
        return res.stream().mapToInt(Integer::intValue).toArray();
    }
}