~/problems / Arrays & hashing / Prefix sums and difference arrays

Count of Stable Subarrays

medium ~25 min CitadelScale AI

Servers stand in a row, and capacity[i] is the processing capacity of server i. The operations team looks for stable stretches: a contiguous block of servers l..r (inclusive) is stable when

  1. it holds at least three servers (r - l >= 2), and
  2. the two end servers have the same capacity, and that capacity equals the total capacity of all the servers strictly between them: capacity[l] == capacity[r] == capacity[l+1] + ... + capacity[r-1].

Write count_stable(capacity: list[int]) -> int that returns how many pairs (l, r) form a stable block.

count_stable([4, 1, 3, 4, 4, 0, 4])   # 2
#   (0, 3): ends 4 and 4, middle 1 + 3 = 4       stable
#   (3, 6): ends 4 and 4, middle 4 + 0 = 4       stable
#   (4, 6): ends 4 and 4, middle 0               not stable
#   (0, 4): ends 4 and 4, middle 1 + 3 + 4 = 8   not stable
count_stable([5, 2, 3, 5])            # 1   the whole row
count_stable([0, 0, 0, 0])            # 3   (0, 2), (1, 3) and (0, 3)
count_stable([7, 7])                  # 0   too short

Constraints: 0 <= len(capacity) <= 200,000 and 0 <= capacity[i] <= 10**9. Checking every pair is O(n²) and too slow at this size.

Show hint

With prefix sums P, the middle of l..r sums to P[r] - P[l + 1]. For a fixed right end r, you need earlier left ends l <= r - 2 with capacity[l] == capacity[r] and P[l + 1] == P[r] - capacity[r]. Count them with a dictionary keyed by that pair.

Topic: Prefix sums and difference arrays. O(1) range sums (1D and 2D); O(1) range updates with difference arrays.

Read the visual guide
0:00
Ctrl ' run · Ctrl ↵ submit
esc