Count stable server segments

A segment [l, r] is stable when it has at least 3 elements, its two end values are equal, and each end value equals the sum of the values strictly between them. In formula terms: r - l >= 2, a[l] == a[r], and a[l] == a[l+1] + ... + a[r-1].

Brute force. Try every pair (l, r) and sum the interior. That costs O(n^3). If you keep a running interior sum while r moves right, it costs O(n^2). Start here in the interview. It pins down the definition and gives you an oracle for testing.

Turn the interior sum into prefix sums. Let P[i] = a[0] + ... + a[i-1], so P[0] = 0. The interior sum is P[r] - P[l+1]. The segment is stable when:

• a[r] == a[l], and
• P[r] == P[l+1] + a[l].

The left endpoint l alone fixes the key (a[l], P[l+1] + a[l]). The right endpoint r alone fixes the key (a[r], P[r]). A pair is stable exactly when the two keys are equal. The value has to be part of the key: two positions with matching prefix sums but different values do not form a stable segment.

One pass with a hash map. Sweep r from left to right and keep a counter of left keys. For each r, first add the left key of l = r - 2. Then add count[(a[r], P[r])] to the answer. Adding l two steps late is what enforces length >= 3. Without it, [0, 0] would count, because its empty interior sums to 0.

Invariant. Before we query at r, the map holds the left keys of exactly the indices 0 .. r-2. These are the valid left endpoints for r, and the map holds nothing else.

Nothing here assumes positive values, so zeros and negative numbers work. The bound on |a[i]| only affects integer size, and Python integers do not overflow. In Java or C++, use 64-bit sums.

from collections import defaultdict
from typing import Dict, List, Tuple


def count_stable_brute(capacity: List[int]) -> int:
    """O(n^2) reference: running interior sum for each left endpoint."""
    n = len(capacity)
    total = 0
    for l in range(n):
        inner = 0
        for r in range(l + 2, n):
            inner += capacity[r - 1]          # interior is capacity[l+1 .. r-1]
            if capacity[l] == capacity[r] == inner:
                total += 1
    return total


def count_stable_subarrays(capacity: List[int]) -> int:
    """Count segments [l, r] with r - l >= 2 and
    capacity[l] == capacity[r] == sum(capacity[l+1:r]). O(n) expected."""
    n = len(capacity)
    if n < 3:
        return 0

    prefix = [0] * (n + 1)                    # prefix[i] = sum(capacity[:i])
    for i, x in enumerate(capacity):
        prefix[i + 1] = prefix[i] + x

    # key -> number of left endpoints l (already eligible) with that key
    left_keys: Dict[Tuple[int, int], int] = defaultdict(int)
    total = 0
    for r in range(2, n):
        l = r - 2                             # becomes eligible now (length >= 3)
        left_keys[(capacity[l], prefix[l + 1] + capacity[l])] += 1
        total += left_keys.get((capacity[r], prefix[r]), 0)
    return total


def list_stable_subarrays(capacity: List[int]) -> List[Tuple[int, int]]:
    """Follow-up: return the (l, r) pairs themselves. O(n + output)."""
    n = len(capacity)
    prefix = [0] * (n + 1)
    for i, x in enumerate(capacity):
        prefix[i + 1] = prefix[i] + x
    lefts: Dict[Tuple[int, int], List[int]] = defaultdict(list)
    out: List[Tuple[int, int]] = []
    for r in range(2, n):
        l = r - 2
        lefts[(capacity[l], prefix[l + 1] + capacity[l])].append(l)
        for left in lefts.get((capacity[r], prefix[r]), []):
            out.append((left, r))
    return out


if __name__ == "__main__":
    import random

    assert count_stable_subarrays([]) == 0
    assert count_stable_subarrays([5, 5]) == 0
    assert count_stable_subarrays([4, 4, 4]) == 1
    assert count_stable_subarrays([6, 1, 2, 3, 6]) == 1
    assert count_stable_subarrays([2, 2, 2, 2, 2]) == 3
    assert count_stable_subarrays([0, 0, 0, 0]) == 3
    assert count_stable_subarrays([7, 3, 4, 7, 2, 5, 7]) == 2
    assert count_stable_subarrays([1, 2, 3]) == 0
    assert count_stable_subarrays([-1, 2, -3, -1]) == 1
    assert list_stable_subarrays([7, 3, 4, 7, 2, 5, 7]) == [(0, 3), (3, 6)]

    for _ in range(3000):
        arr = [random.randint(-3, 4) for _ in range(random.randint(0, 12))]
        assert count_stable_subarrays(arr) == count_stable_brute(arr), arr
        assert len(list_stable_subarrays(arr)) == count_stable_brute(arr), arr
    print("ok")

• count_stable_subarrays: O(n) expected time for the prefix pass and the sweep, with O(1) average per hash-map operation. O(n) space for the prefix array and the map. You can drop the prefix array by keeping a running sum, but the map stays O(n).
• list_stable_subarrays: O(n + k) time, where k is the number of stable segments. k can reach Θ(n^2), for example on an all-zero array.
• count_stable_brute: O(n^2) time, O(1) extra space.

• [], [5], [5, 5] → 0. Fewer than 3 elements is never stable.
• [4, 4, 4] → 1. The interior is [4].
• [1, 2, 3] → 0. The ends differ.
• [6, 1, 2, 3, 6] → 1. 1 + 2 + 3 = 6.
• [2, 2, 2, 2, 2] → 3. Each length-3 window qualifies, and no longer window does.
• [0, 0, 0, 0] → 3. [0..2], [1..3] and [0..3] all qualify. [0, 0] of length 2 must not count.
• [7, 3, 4, 7, 2, 5, 7] → 2. The stable segments are [0..3] (3 + 4 = 7) and [3..6] (2 + 5 = 7). The outer segment [0..6] has equal ends, but its interior sums to 21, so it is not stable.
• Negative values: [-1, 2, -3, -1] → 1. The interior sums to -1.
• Randomized check: the linear version matches the brute force on thousands of small random arrays with values in [-3, 4].

• Return the segments, not just the count. Store a list of left indices per key instead of a counter (list_stable_subarrays). Say up front that the output can be quadratic, for example on an all-zero array.
• Streaming input. Capacities arrive one at a time. Keep the running prefix sum, the last two values and the map, and update the count in O(1) per arrival. Memory is O(distinct keys). If only segments of length up to W matter, evict keys older than W using a queue.
• Change the stability rule. Examples: the interior sum equals k * a[l], the ends differ by at most d, or there is no length minimum. The pattern stays the same: split the condition into a function of l and a function of r, then match the keys in a hash map. An exact equality keeps O(n). A tolerance such as "within d" needs a sorted structure or a Fenwick tree over compressed keys, which costs O(n log n).
• Range queries. Answer "how many stable segments lie inside [L, R]" for many queries. Process them offline: sort by R, sweep, and use Mo's algorithm, or a Fenwick tree over left endpoints keyed by the matching key. Expect O((n + q) sqrt n) or O((n + q) log n).
• Point updates. Changing one capacity shifts every prefix sum after it, so the map keys move. Discuss sqrt decomposition, or answer that a full rebuild is O(n) per update and is acceptable for rare updates.

Mid-level (L4)

• States the definition precisely, including length >= 3 and both ends equal, and writes a correct O(n^2) solution.
• With a hint, moves to prefix sums and a hash map, and gets the length-3 offset right.
• Tests the small and degenerate cases by hand.

Senior (L5)

• Arrives at the (value, prefix) key matching without help and explains why the value must be part of the key.
• Delays insertion by two to enforce the minimum length, and catches the all-zero trap.
• Gives clean complexity, mentions 64-bit sums in fixed-width languages, and checks against the brute force.

Staff (L6+)

• Frames the method as a general technique: separate the condition into f(l) == g(r) and match with a hash map. Applies it quickly to variant rules.
• Handles the streaming, range-query and enumeration follow-ups, including the quadratic output bound and the tradeoffs between the offline and online approaches.
• Writes a randomized differential test without being asked.