Count redundant substrings by vowel and consonant weights

Question: Count redundant substrings by vowel and consonant weights

Companies: Amazon

Difficulty: Medium

Problem statement

Write a function that counts the redundant substrings of a word.

A string s is redundant when its length equals a * V + b * C, where V is the number of vowels in s, C is the number of consonants in s, and a and b are given integers. The vowels are a, e, i, o, u; every other letter (including y) is a consonant.

def count_redundant_substrings(word: str, a: int, b: int) -> int: ...

Return how many of the n(n+1)/2 non-empty contiguous substrings of word are redundant. Substrings at different positions count separately, even if their text is the same.

Constraints

• 0 <= len(word) <= 10^5, lowercase English letters only.
• -10^5 <= a, b <= 10^5.
• The answer can exceed 2^31.

Example 1

word = "code", a = 2, b = 0  ->  4

The rule becomes V + C = 2V, so V == C. The redundant substrings are "co", "od", "de" and "code".

Example 2

word = "queue", a = 1, b = 0  ->  10

The rule becomes V + C = V, so C == 0. Only substrings made entirely of vowels count, and "ueue" has 10 of them.

An O(n^2) solution is too slow for the largest inputs; aim for O(n).

Approach

A substring has length V + C, where V counts its vowels and C counts its consonants. It is redundant when

V + C = a*V + b*C

Move everything to one side:

(1 - a)*V + (1 - b)*C = 0

So give every character a weight: a vowel weighs 1 - a, a consonant weighs 1 - b. A substring is redundant exactly when the sum of its weights is zero. The question becomes "count the subarrays with sum zero", a standard prefix-sum problem.

Brute force. Try every start i. Extend the end j one character at a time and keep running counts of vowels and consonants. Check j - i + 1 == a*V + b*C at each step. That is O(n^2) time and O(1) space. It is correct and worth saying first, but with n up to about 10^5 it times out.

Optimal: prefix sums and a hash map.

• Let P[0] = 0 and P[k] = weight(word[0]) + ... + weight(word[k-1]).
• The substring word[i:j] has weight P[j] - P[i]. It is redundant when P[j] == P[i] with i < j.
• Walk left to right. Keep a counter of how many times each prefix value has appeared so far. At each new prefix value p, every earlier equal prefix closes one redundant substring, so add seen[p] to the answer, then increment seen[p].

Each pair i < j is counted exactly once, when the scan reaches j. Seeding seen[0] = 1 (for P[0]) counts substrings that start at index 0.

Why it is correct. The equation is linear in V and C, and both counts are additive over concatenation. That is what lets one integer per prefix stand in for the pair (V, C). If the condition were not linear (for example V * C == length), this trick would not apply.

Special cases fall out for free.

• a = 1 and b = 1: both weights are 0, every prefix is 0, and the answer is n(n+1)/2. Every substring is redundant.
• a = 1 only: vowels weigh 0, so a substring is redundant exactly when it has no consonants.
• Weights can be negative or zero; the hash map does not care.

Invariants.

• After processing k characters, prefix == P[k] and seen[x] is the number of indices t <= k with P[t] == x.
• count equals the number of redundant substrings that end at or before index k - 1.

Code

from collections import defaultdict
from typing import Dict

VOWELS = frozenset("aeiou")


def count_redundant_substrings(word: str, a: int, b: int) -> int:
    """Count substrings s of word with len(s) == a * vowels(s) + b * consonants(s).

    word holds lowercase English letters. 'y' counts as a consonant.
    """
    vowel_weight = 1 - a
    consonant_weight = 1 - b

    seen: Dict[int, int] = defaultdict(int)
    seen[0] = 1  # the empty prefix P[0]
    prefix = 0
    count = 0
    for ch in word:
        prefix += vowel_weight if ch in VOWELS else consonant_weight
        count += seen[prefix]
        seen[prefix] += 1
    return count


def count_redundant_brute(word: str, a: int, b: int) -> int:
    """O(n^2) reference: extend every start and test the length rule."""
    n = len(word)
    total = 0
    for i in range(n):
        v = c = 0
        for j in range(i, n):
            if word[j] in VOWELS:
                v += 1
            else:
                c += 1
            if j - i + 1 == a * v + b * c:
                total += 1
    return total


if __name__ == "__main__":
    import random

    # Equal vowels and consonants: "co", "od", "de", "code".
    assert count_redundant_substrings("code", 2, 0) == 4
    # a = b = 1: every substring counts.
    assert count_redundant_substrings("rhythm", 1, 1) == 21
    # a = 1, b = 0: substrings made only of vowels.
    assert count_redundant_substrings("queue", 1, 0) == 10
    assert count_redundant_substrings("", 3, -2) == 0
    assert count_redundant_substrings("z", 0, 1) == 1
    assert count_redundant_substrings("z", 0, 0) == 0

    rng = random.Random(7)
    for _ in range(2000):
        w = "".join(rng.choice("aeioubcdxyz") for _ in range(rng.randint(0, 12)))
        a, b = rng.randint(-3, 3), rng.randint(-3, 3)
        assert count_redundant_substrings(w, a, b) == count_redundant_brute(w, a, b)
    print("all tests passed")

Complexity

• Optimal: O(n) time for one pass with O(1) hash map operations. O(n) space in the worst case for the map of distinct prefix values. When |1 - a| and |1 - b| are small, prefix values lie in [-n*w, n*w] and a plain array offset by n*w can replace the hash map.
• Brute force: O(n^2) time, O(1) space.

Test cases

• word = "code", a = 2, b = 0 → 4. Weights are -1 for vowels and +1 for consonants, so a substring counts when its vowels equal its consonants: "co", "od", "de", "code".
• word = "rhythm", a = 1, b = 1 → 21. All weights are 0; every one of the 6*7/2 substrings counts.
• word = "queue", a = 1, b = 0 → 10. Vowels weigh 0 and consonants weigh 1, so only all-vowel substrings count. The vowel run "ueue" has length 4, giving 4*5/2 = 10.
• word = "", a = 3, b = -2 → 0.
• word = "z", a = 0, b = 1 → 1. The consonant weighs 0, so the single-letter substring counts.
• word = "z", a = 0, b = 0 → 0. Length 1 never equals 0.
• word = "aab", a = 3, b = -1 → 1. Vowels weigh -2, consonants weigh 2; only "ab" sums to 0.
• Large input, n = 10^5 of one letter with a = b = 1 → 5000050000. Checks that the count does not overflow in languages with 32-bit ints.

Follow-ups

1. Return the substrings, or the longest one. Store the first index of each prefix value instead of a count. The longest redundant substring ending at j starts at first[P[j]]. Listing all of them is O(n + output) by keeping each prefix value's index list.
2. Answer many (a, b) queries on the same word. Each query needs its own weights, so it is O(n) per query. If the queries share a and b only through the ratio (1 - a) : (1 - b), scale the weights by their gcd and reuse one pass per distinct ratio.
3. Streaming input. The algorithm is already one pass. Keep prefix, count and seen as state and report count after each character. Memory stays O(distinct prefixes).
4. Generalize the rule to length = a*V + b*C + k. Rewrite it as sum(weights) = k. Look up seen[prefix - k] instead of seen[prefix], the same trick as "subarray sum equals k".
5. Count only substrings of length at least L. Delay inserting P[t] into seen until the scan reaches index t + L, so only starts far enough back are counted.

Interviewer expectations

Mid-level (L4)

• Writes the O(n^2) brute force correctly, then finds the algebra that turns the rule into a zero-sum condition.
• Implements the prefix-sum counter with the seen[0] = 1 seed and explains why it is needed.
• Tests the empty string and the a = b = 1 case.

Senior (L5)

• Goes straight to the reduction and states why linearity makes it valid.
• Handles negative and zero weights and calls out the 64-bit count for large inputs.
• Adapts quickly to "sum equals k" and "longest substring" variants.

Staff (L6+)

• Frames the problem as a family (weighted zero-sum subarrays) and says which rule changes would break the reduction.
• Discusses array versus hash map storage, memory for streaming input, and batching many (a, b) queries.
• Writes a brute-force cross-check or property test without being asked.