Split a string into blocks that match the prefix's letter counts: Frontier Interview's solution

Question: Split a string into blocks that match the prefix's letter counts

Companies: Amazon

Difficulty: Easy

Problem statement

You get a string s of uppercase letters A-Z.

Call a positive length L valid when:

L <= len(s) and s can be cut into consecutive, non-overlapping blocks of exactly L characters with nothing left over, and
every block contains each letter the same number of times as the prefix s[:L] does.

The blocks do not have to be equal strings; only their letter counts must match the prefix.

Implement valid_prefix_lengths(s: str) -> list[int] that returns all valid lengths in increasing order. (len(s) is always valid.)

Constraints: 1 <= len(s) <= 10^5.

Example 1

s = "ABCCBAACB"
output: [3, 9]

With L = 3 the blocks are ABC, CBA, ACB, each with one A, one B and one C. L = 1 fails because B differs from A.

Example 2

s = "ABBAAABB"
output: [4, 8]

L = 2 fails: the prefix AB has one of each letter, but the block AA does not. L = 4 works: ABBA and AABB each have two As and two Bs.

Reconstructed from the reported question; details may differ from what interviewers ask.

Approach

A length L can only work if the string splits evenly, so L must divide
n = len(s). That cuts the candidates from n to the divisors of n
(at most a few hundred even for n = 10^6).

Brute force. For every L from 1 to n with n % L == 0, build a
Counter for each block and compare it to the prefix's counter. Each check
touches every character once, so one candidate costs O(n) and all
candidates cost O(n * d(n)), where d(n) is the number of divisors. That
is already acceptable for an "Easy" round, and it is the right thing to say
first.

Better: prefix counts. Build pref[i][c], the number of times letter c
appears in s[:i]. The count of c in any block s[a:b] is then
pref[b][c] - pref[a][c], an O(26) lookup. Checking length L costs
O(26 * n / L), and summing over all divisors gives O(26 * sigma(n)),
where sigma(n) is the sum of divisors of n (O(n log log n)). Two cheap
filters reject most candidates before any block is examined:

L must divide n.
For every letter, the total count in s must equal (n / L) times its
count in the prefix. If not, some block must differ from the prefix.

Filter 2 is necessary but not sufficient. In "ABABBA" with L = 3, the
totals (A=3, B=3) are not twice the prefix "ABA" (A=2, B=1), so it is
rejected early. But "ABBAAABB" with L = 2 passes the totals
(A=4, B=4, four times the prefix "AB") while the blocks "AA" and "BB"
differ. So the per-block pass is still required.

Invariants.

L = n is always valid: the only block is the prefix itself.
If L is valid and M is a multiple of L that divides n, then M is
valid too: each M-block is the union of M / L consecutive L-blocks,
and so is the M-prefix, so they all have M / L times the base counts.

The function returns every valid length in increasing order; the smallest
one is result[0].

Code

from typing import List


def _divisors(n: int) -> List[int]:
    small, large = [], []
    d = 1
    while d * d <= n:
        if n % d == 0:
            small.append(d)
            if d != n // d:
                large.append(n // d)
        d += 1
    return small + large[::-1]


def valid_prefix_lengths(s: str) -> List[int]:
    """Return every prefix length L such that s splits into blocks of
    length L, each with the same letter counts as s[:L]. Sorted ascending."""
    n = len(s)
    if n == 0:
        return []

    # pref[i][c] = occurrences of letter c in s[:i]
    pref: List[List[int]] = [[0] * 26]
    for ch in s:
        row = pref[-1][:]
        row[ord(ch) - 65] += 1
        pref.append(row)

    total = pref[n]
    result: List[int] = []
    for length in _divisors(n):
        blocks = n // length
        target = pref[length]

        # Cheap filter: totals must be `blocks` copies of the prefix.
        if any(total[c] != blocks * target[c] for c in range(26)):
            continue

        ok = True
        for start in range(length, n, length):
            lo, hi = pref[start], pref[start + length]
            if any(hi[c] - lo[c] != target[c] for c in range(26)):
                ok = False
                break
        if ok:
            result.append(length)
    return result


def smallest_valid_length(s: str) -> int:
    lengths = valid_prefix_lengths(s)
    return lengths[0] if lengths else 0


if __name__ == "__main__":
    assert valid_prefix_lengths("ABBA") == [2, 4]
    assert valid_prefix_lengths("ABCCBAACB") == [3, 9]
    assert valid_prefix_lengths("AABB") == [4]
    assert valid_prefix_lengths("ZZZZ") == [1, 2, 4]
    assert valid_prefix_lengths("Q") == [1]
    assert valid_prefix_lengths("ABABBA") == [2, 6]
    assert valid_prefix_lengths("ABBAAB") == [2, 6]
    assert valid_prefix_lengths("ABAB" * 3) == [2, 4, 6, 12]
    assert valid_prefix_lengths("") == []
    assert smallest_valid_length("BAABAB") == 2
    print("ok")

Complexity

Prefix table: O(26 n) time and space. To save memory, keep one running
Counter per candidate instead; then space is O(26) per candidate and
time is O(n) per candidate.
Divisor enumeration: O(sqrt n).
Checking all candidates: sum over divisors L of 26 * (n / L) =
O(26 * sigma(n)), which is O(26 * n log log n) in the worst case and
near O(26 n) in practice.
Overall: O(26 n log log n) time, O(26 n) space.

Test cases

"ABBA" → [2, 4] (blocks "AB", "BA" share counts A=1, B=1).
"ABCCBAACB" → [3, 9] ("ABC", "CBA", "ACB").
"AABB" → [4] (L = 2 gives "AA" vs "BB"; L = 1 gives "A" vs "B").
"ZZZZ" → [1, 2, 4] (every divisor works for a single repeated letter).
"Q" → [1] (single character).
"ABABBA" → [2, 6] (L = 3: "ABA" vs "BBA" fails).
"ABAB" * 3 (length 12) → [2, 4, 6, 12] (L = 3: "ABA" vs "BAB" fails).
"" → [] (no positive prefix length; state the convention with the interviewer).
"ABBAAABB" → [4, 8] (passes the totals filter at L = 2 but fails on "AA").
"BAABAB" smallest valid length → 2.

Follow-ups

Return only the count or only the smallest valid length. Same loop;
stop at the first success for the smallest. Mention the monotonic
property: once L is valid, any multiple of L that divides n is valid,
so you can mark those without checking.
Answer for every prefix of s, not just the whole string (the
"maximum equal parts per prefix" variant). Now the string length changes
with each query. Maintain running counts; for prefix length i, test
only divisors of i. Precompute divisor lists with a sieve in
O(n log n). Warn the interviewer that "GCD of the letter counts" is a
necessary condition, not a sufficient one, when blocks must be contiguous:
"AABB" has GCD 2, yet "AA" and "BB" differ.
Streaming input. Characters arrive one at a time and you must report
whether the current string has a valid length shorter than itself. Keep
the prefix counts table incrementally (or a rolling hash of the count
vector per block) and re-check divisors of the new length only.
Larger alphabet (Unicode). Replace the fixed 26-slot arrays with
dictionaries keyed only by letters present in the prefix; compare a block
by its own dictionary. Or hash count vectors (sum of random weights per
letter) to compare blocks in O(1), with a full check on hash match.
Blocks need not be contiguous in order / allow one mismatch. For "at
most k blocks may differ", count failing blocks instead of breaking on the
first one.

Interviewer expectations

Mid-level (L4)

Restates the rule correctly: blocks must match counts, not be identical.
Notices that L must divide n and that L = n is always valid.
Writes a working Counter-per-block solution and tests "AABB" and a
single-letter string.

Senior (L5)

Uses prefix count arrays so each block check is O(26).
Adds the totals filter and the early break, and states the
divisor-sum complexity precisely.
Handles empty input and explains the "multiple of a valid length is
valid" property.

Staff (L6+)

Extends cleanly to the per-prefix and streaming variants without
rewriting the core check.
Discusses why the GCD-of-counts shortcut is wrong for contiguous blocks,
with a counter-example.
Talks about memory for very long strings (26 x n table vs running
counters) and hashing count vectors for large alphabets.