Power Grid Maintenance

main.md

Question

Approach

I first use a Disjoint Set Union (Union-Find) to determine the connected component (power grid) of each station. For every component root I build a min-heap of all station ids in that component (initially all stations are online). I keep an online boolean array. When a station goes offline I simply mark online[id] = False. For a query [1, x]:

• if x is online I return x,
• otherwise I look up the heap of x's component and pop offline ids from the top until the heap is empty or the top is an online station; if the heap is empty I return -1, else I return the heap top (the smallest online id in that component).

This avoids expensive deletes from heaps. Each id is popped at most once over all queries.

Implementation

class DSU:
    def __init__(self, n: int):
        self.parent = list(range(n+1))
        self.rank = [0]*(n+1)
    def find(self, x: int) -> int:
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]
    def union(self, a: int, b: int) -> None:
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return
        if self.rank[ra] < self.rank[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra
        if self.rank[ra] == self.rank[rb]:
            self.rank[ra] += 1

class Solution:
    def processQueries(self, c: int, connections: List[List[int]], queries: List[List[int]]) -> List[int]:
        dsu = DSU(c)
        for u, v in connections:
            dsu.union(u, v)
        heaps = {}
        for node in range(1, c+1):
            root = dsu.find(node)
            if root not in heaps:
                heaps[root] = []
            heaps[root].append(node)
        for root in heaps:
            heapq.heapify(heaps[root])

        online = [True] * (c + 1) 
        res = []
        for typ, x in queries:
            if typ == 2:
                if 1 <= x <= c:
                    online[x] = False
            else:
                if online[x]:
                    res.append(x)
                else:
                    root = dsu.find(x)
                    heap = heaps.get(root, [])
                    while heap and not online[heap[0]]:
                        heapq.heappop(heap)
                    if not heap:
                        res.append(-1)
                    else:
                        res.append(heap[0])
        return res

Complexities

• Time: O(1)
• Space: O(c)