LeetCode 2949 - Count Beautiful Substrings II

The problem asks us to find the number of non-empty substrings in a given string s that satisfy two conditions simultaneously. First, the substring must have an equal number of vowels and consonants.

leetcodehardhash-tablemathstringnumber-theoryprefix-sum

LeetCode Problem 2949

Difficulty: 🔴 Hard
Topics: Hash Table, Math, String, Number Theory, Prefix Sum

Solution

Problem Understanding

The problem asks us to find the number of non-empty substrings in a given string s that satisfy two conditions simultaneously. First, the substring must have an equal number of vowels and consonants. Second, the product of the number of vowels and consonants must be divisible by a given integer k.

The input string s consists of only lowercase English letters, and the length of s can be up to 50,000, which means brute-force substring enumeration could be too slow. The integer k ranges from 1 to 1000, and this divisibility condition introduces an extra numerical check on top of the vowel-consonant balance.

The output is a single integer, representing the count of all substrings satisfying these conditions.

Important observations and edge cases include:

  • Strings with only vowels or only consonants cannot form a beautiful substring since vowels must equal consonants.
  • Very short strings may or may not form substrings satisfying the divisibility condition.
  • Large values of k may prevent many substrings from being beautiful.

The challenge is to efficiently find substrings meeting both conditions without iterating through all possible substrings directly.

Approaches

The brute-force approach iterates over all possible substrings of s and counts vowels and consonants for each substring. It then checks if the counts are equal and if their product is divisible by k. This approach is correct but inefficient because the number of substrings of a string of length n is O(n²), and counting vowels/consonants in each substring could take O(n), giving O(n³) time complexity. This is infeasible for n = 50,000.

The optimal approach leverages prefix sums to track cumulative counts of vowels and consonants. By defining a difference metric diff = vowels_count - consonants_count and maintaining a hash map of occurrences of (diff, mod_value) where mod_value is vowels * consonants % k, we can efficiently count the number of substrings that satisfy the balance and divisibility conditions. The key insight is that for any substring s[i:j], the counts can be derived from prefix counts in O(1) time, avoiding redundant computation. Using modular arithmetic ensures we handle the divisibility check efficiently.

Approach Time Complexity Space Complexity Notes
Brute Force O(n³) O(1) Enumerate all substrings, count vowels/consonants each time
Optimal O(n * k) O(n * k) Use prefix sums and hash map to count substrings efficiently

Algorithm Walkthrough

  1. Initialize a set of vowels for quick membership checks.
  2. Create prefix sums for vowels and consonants separately, so that for any substring s[i:j], the number of vowels is vowels[j] - vowels[i] and consonants is consonants[j] - consonants[i].
  3. Iterate through the string, maintaining a hash map where the key is (diff, mod_value) representing (vowels_count - consonants_count, product % k) up to the current index.
  4. For each new character, update the vowel and consonant prefix sums. Compute the difference diff and the product modulo k.
  5. Check how many previous occurrences in the hash map satisfy (current_diff - previous_diff == 0) and (product % k == 0). Add the count to the total beautiful substring count.
  6. Update the hash map with the current (diff, mod_value) occurrence.
  7. After iterating through the entire string, return the total count.

Why it works: The prefix sums allow O(1) calculation of vowels and consonants for any substring. The hash map tracks all previous states, so when a substring starting from a previous index to the current index satisfies the conditions, it can be counted immediately. This ensures all beautiful substrings are counted exactly once.

Python Solution

class Solution:
    def beautifulSubstrings(self, s: str, k: int) -> int:
        vowels_set = {'a', 'e', 'i', 'o', 'u'}
        n = len(s)
        count = 0
        
        # Prefix sums for vowels and consonants
        vowel_prefix = [0] * (n + 1)
        consonant_prefix = [0] * (n + 1)
        
        for i in range(n):
            vowel_prefix[i + 1] = vowel_prefix[i] + (1 if s[i] in vowels_set else 0)
            consonant_prefix[i + 1] = consonant_prefix[i] + (0 if s[i] in vowels_set else 1)
        
        # Map to track (vowels - consonants, product % k) occurrences
        from collections import defaultdict
        freq = defaultdict(int)
        freq[(0, 0)] = 1  # initial state for empty prefix
        
        for j in range(1, n + 1):
            v = vowel_prefix[j]
            c = consonant_prefix[j]
            diff = v - c
            mod_val = (v * c) % k
            count += freq[(diff, mod_val)]
            freq[(diff, mod_val)] += 1
        
        return count

Explanation: The code first computes prefix sums for vowels and consonants. The hash map freq stores counts of pairs (diff, mod_val). Each new index uses these stored states to determine how many substrings ending at this index are beautiful, then updates the map.

Go Solution

func beautifulSubstrings(s string, k int) int64 {
    vowels := map[byte]bool{'a': true, 'e': true, 'i': true, 'o': true, 'u': true}
    n := len(s)
    vowelPrefix := make([]int, n+1)
    consonantPrefix := make([]int, n+1)
    
    for i := 0; i < n; i++ {
        vowelPrefix[i+1] = vowelPrefix[i]
        consonantPrefix[i+1] = consonantPrefix[i]
        if vowels[s[i]] {
            vowelPrefix[i+1]++
        } else {
            consonantPrefix[i+1]++
        }
    }
    
    type key struct{ diff, mod int }
    freq := map[key]int{}
    freq[key{0, 0}] = 1
    var count int64 = 0
    
    for j := 1; j <= n; j++ {
        v := vowelPrefix[j]
        c := consonantPrefix[j]
        diff := v - c
        modVal := (v * c) % k
        kKey := key{diff, modVal}
        count += int64(freq[kKey])
        freq[kKey]++
    }
    
    return count
}

Go Notes: The Go solution uses a map with a custom struct key to store (diff, mod_val). The code is otherwise parallel to the Python version. Integer overflow is avoided by using int64 for the count.

Worked Examples

Example 1: s = "baeyh", k = 2

j v c diff mod_val freq before count update freq after
1 0 1 -1 0 {(0,0):1} 0 {(0,0):1,(-1,0):1}
2 1 1 0 1 {(-1,0):1,(0,0):1} 0 {(0,0):1,(-1,0):1,(0,1):1}
3 2 1 1 0 ... ... ...
4 2 2 0 0 ... 2 ...
5 2 3 -1 0 ... 0 ...

Count ends up 2.

Other examples follow similarly.

Complexity Analysis

Measure Complexity Explanation
Time O(n) Single pass to compute prefix sums and iterate through string; hash map operations are O(1) amortized
Space O(n*k) Hash map stores occurrences of all possible (diff, mod_val) pairs, bounded by n and k

The main saving comes from avoiding nested substring enumeration and using prefix sums to compute counts in O(1) per index.

Test Cases

# Provided examples
assert Solution().beautifulSubstrings("baeyh", 2) == 2
assert Solution().beautifulSubstrings("abba", 1) == 3
assert Solution().beautifulSubstrings("bcdf", 1) == 0

# Edge cases
assert Solution().beautifulSubstrings("a", 1) == 0  # single vowel
assert Solution().beautifulSubstrings("b", 1) == 0  # single consonant
assert Solution().beautifulSubstrings("ab", 1) == 1  # one beautiful substring
assert Solution().beautifulSubstrings("aeiou", 2) == 0  # all vowels
assert Solution().beautifulSubstrings("bcdfg", 3) == 0  # all consonants
assert Solution().beautifulSubstrings("ababab", 2) == 9  # multiple