Skip to main content
Back to problems
#3607
Medium Algorithms

Power grid maintenance

Array Hash Table Depth-First Search Breadth-First Search Union-Find Graph Theory Heap (Priority Queue) Ordered Set
56.3% acceptance
Feb 25, 2026
526
57
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: [1, x]: A maintenance check is requested for station x. If station x is online, it resolves the check by itself. If station x is offline, the check is resolved by the operational station with the smallest id in the same power grid as x. If no operational station exists in that grid, return -1. [2, x]: Station x goes offline (i.e., it becomes non-operational). 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.

Solution

Rust
Time O(n log n)
Space O(n)
LeetCode
solution.rs
impl Solution {
  pub fn process_queries(c: i32, connections: Vec<Vec<i32>>, queries: Vec<Vec<i32>>) -> Vec<i32> {
    use std::collections::BTreeSet;
    let c = c as usize;
    let mut parent = (0..c).collect::<Vec<_>>();
    let mut rank = vec![0usize; c];
    // Each root owns a BTreeSet of active station ids (1-indexed)
    let mut sets: Vec<BTreeSet<i32>> = (0..c).map(|i| {
      let mut s = BTreeSet::new();
      s.insert(i as i32 + 1);
      s
    }).collect();
    let mut online = vec![true; c + 1]; // 1-indexed

    fn find(parent: &mut Vec<usize>, x: usize) -> usize {
      if parent[x] != x { parent[x] = find(parent, parent[x]); }
      parent[x]
    }

    // Build DSU from connections
    for conn in &connections {
      let (u, v) = (conn[0] as usize - 1, conn[1] as usize - 1);
      let ru = find(&mut parent, u);
      let rv = find(&mut parent, v);
      if ru == rv { continue; }
      // Union by rank, merge smaller set into larger
      let (keep, give) = if rank[ru] >= rank[rv] { (ru, rv) } else { (rv, ru) };
      if rank[keep] == rank[give] { rank[keep] += 1; }
      parent[give] = keep;
      let give_set = std::mem::take(&mut sets[give]);
      sets[keep].extend(give_set);
    }

    let mut result = Vec::new();
    for q in &queries {
      let (t, x) = (q[0], q[1] as usize);
      if t == 2 {
        online[x] = false;
        let root = find(&mut parent, x - 1);
        sets[root].remove(&(x as i32));
      } else {
        if online[x] {
          result.push(x as i32);
        } else {
          let root = find(&mut parent, x - 1);
          match sets[root].iter().next() {
            Some(&v) => result.push(v),
            None     => result.push(-1),
          }
        }
      }
    }
    result
  }
}