We use cookies (including Google cookies) to personalize ads and analyze traffic. By continuing to use our site, you accept our Privacy Policy.

Count Beautiful Substrings II

Number: 3208

Difficulty: Hard

Paid? No

Companies: Amazon


Problem Description

Given a lowercase English string s and a positive integer k, count the number of non‐empty “beautiful” substrings. A substring is called beautiful if:

  • The number of vowels equals the number of consonants.
  • The product (vowels * consonants) is divisible by k. Remember that vowels are "a", "e", "i", "o", "u" and consonants are all other lowercase letters. A substring is a contiguous block within s.

Key Insights

  • Because vowels must equal consonants, any beautiful substring has even length.
  • If vowels = consonants = t, then the product equals t²; hence the condition becomes t² % k == 0.
  • For a given k the condition t² % k == 0 is equivalent to requiring t to be a multiple of L, where L is computed from the prime factorization of k (specifically, for every prime p^a dividing k, t must be divisible by p^(ceil(a/2))).
  • We can use prefix sums to count balanced substrings. Define two sequences:
    • diff[i]: difference between number of vowels and consonants from the beginning up to index i.
    • vowelCount[i]: number of vowels up to index i.
  • A substring s[i...j-1] is balanced (vowels == consonants) if diff[j] == diff[i]. Its vowel count is vowelCount[j] – vowelCount[i] and must be positive and a multiple of L.
  • Since prefix vowel counts are non‐decreasing, grouping indices with the same diff allows us to reduce the problem to “count pairs (i, j) with i < j” in which the difference (vowelCount[j] – vowelCount[i]) is a positive multiple of L.
  • For any group with fixed diff, note that for indices i and j their contribution is valid only if their difference in vowel counts is non zero and (vowelCount[j] – vowelCount[i]) % L == 0.
  • By partitioning each diff-group further by the residue (vowelCount mod L), the condition becomes: within each residue class, count only those pairs for which the vowel counts are not equal (since equal vowel count gives difference 0).
  • For each residue class, if there are p prefix indices in total and, among these, if a specific vowel count appears f times, then the total ordered pairs in that residue class is (p choose 2) and the invalid pairs (with difference 0) is the sum of (f choose 2) over each distinct vowel count. Their difference gives the count of valid pairs for that residue class.

Space and Time Complexity

Time Complexity: O(n + G) where n is the length of s and G is the total number of groups and residue classes; in worst-case (when grouping by diff and then by modulo) it is near O(n) on average. Space Complexity: O(n) for storing prefix arrays and the auxiliary dictionaries.


Solution

The solution uses the following steps:

  1. Precompute L – the smallest number such that for every prime factor p^a of k, L includes p^(ceil(a/2)). This ensures that for any t, t² is divisible by k if and only if t is a multiple of L.
  2. Compute prefix arrays:
    • diff[i] = (# vowels up to i) - (# consonants up to i)
    • vowelCount[i] = (# vowels up to i) (Consider the empty prefix at index 0.)
  3. Group all prefix indices by their diff value. For every group, further group prefix indices by (vowelCount mod L) because a valid difference (vowelCount[j] – vowelCount[i]) is divisible by L if and only if vowelCount[j] and vowelCount[i] have the same remainder modulo L.
  4. In each modulo group, count all pairs (i, j) (with i < j). Because the prefix vowel counts are non-decreasing over the original order, a pair with equal vowel counts means zero difference (invalid). Thus, use counting:
    • Total pairs from the group = (total count choose 2)
    • Remove the pairs that come from indices having the same vowel count (since difference would be 0). For every distinct vowel count in the modulo group, subtract (frequency choose 2).
  5. Sum the counts over all groups and modulo classes to obtain the answer.

Code Solutions

# Python solution for Count Beautiful Substrings II

import math
from collections import defaultdict, Counter

def countBeautifulSubstrings(s, k):
    # Helper function to compute L:
    # For prime factorization of k, for each prime p^e,
    # we require a factor p^(ceil(e/2)) in L.
    def computeL(k):
        L = 1
        # factorize k; since k <= 1000, trial division is acceptable.
        temp = k
        p = 2
        while p*p <= temp:
            if temp % p == 0:
                exp = 0
                while temp % p == 0:
                    exp += 1
                    temp //= p
                L *= p ** math.ceil(exp / 2)
            else:
                p += 1
        if temp > 1:
            # temp is prime
            L *= temp ** math.ceil(1 / 2)
        return L

    L = computeL(k)
    
    n = len(s)
    # Define set of vowels
    vowels = set("aeiou")
    
    # Initialize prefix arrays; index 0 corresponds to empty prefix.
    diff_prefix = [0] * (n + 1)  # vowels count - consonants count
    vowel_prefix = [0] * (n + 1) # vowels count
    for i in range(n):
        # Check if current character is vowel or consonant
        if s[i] in vowels:
            vowel_prefix[i+1] = vowel_prefix[i] + 1
            diff_prefix[i+1] = diff_prefix[i] + 1
        else:
            vowel_prefix[i+1] = vowel_prefix[i]
            diff_prefix[i+1] = diff_prefix[i] - 1

    # Group indices by diff value.
    groups = defaultdict(list)
    for idx, diff_val in enumerate(diff_prefix):
        groups[diff_val].append(idx)

    result = 0
    # For each diff group, further group by (vowel_prefix mod L)
    for diff_val, indices in groups.items():
        # Create dictionary mapping mod value to list of vowel_prefix values.
        mod_groups = defaultdict(list)
        for idx in indices:
            mod_val = vowel_prefix[idx] % L
            mod_groups[mod_val].append(vowel_prefix[idx])
        
        # Process each modulo group
        for mod_val, values in mod_groups.items():
            # Total number of ways to pick i and j (i < j) from this group.
            total_pairs = len(values) * (len(values) - 1) // 2
            # Count invalid pairs where vowel_prefix values are equal (difference 0).
            freq = Counter(values)
            invalid = 0
            for count in freq.values():
                invalid += count * (count - 1) // 2
            # Valid pairs are those with a positive difference that is also a multiple of L.
            result += total_pairs - invalid
    return result

# Example usage:
print(countBeautifulSubstrings("baeyh", 2))  # Expected output: 2
print(countBeautifulSubstrings("abba", 1))   # Expected output: 3
print(countBeautifulSubstrings("bcdf", 1))   # Expected output: 0
← Back to All Questions