Skip to content

2949. Count Beautiful Substrings II 👍

  • Time: $O(|\texttt{s}| + k)$
  • Space: $O(|\texttt{s}| + k)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
class Solution {
 public:
  // Same as 2947. Count Beautiful Substrings I
  int beautifulSubstrings(string s, int k) {
    const int root = getRoot(k);
    int ans = 0;
    int vowels = 0;
    int vowelsMinusConsonants = 0;
    // {(vowels, vowelsMinusConsonants): count}
    unordered_map<pair<int, int>, int, PairHash> prefixCount{{{0, 0}, 1}};

    for (const char c : s) {
      if (isVowel(c)) {
        vowels = (vowels + 1) % root;
        ++vowelsMinusConsonants;
      } else {
        --vowelsMinusConsonants;
      }
      const pair<int, int> prefix{vowels, vowelsMinusConsonants};
      ans += prefixCount[prefix]++;
    }

    return ans;
  }

 private:
  bool isVowel(char c) {
    static constexpr string_view kVowels = "aeiouAEIOU";
    return kVowels.find(c) != string_view::npos;
  }

  int getRoot(int k) {
    for (int i = 1; i <= k; ++i)
      if (i * i % k == 0)
        return i;
    throw;
  }

  struct PairHash {
    size_t operator()(const pair<int, int>& p) const {
      return p.first ^ p.second;
    }
  };
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class Solution {
  // Same as 2947. Count Beautiful Substrings I
  public int beautifulSubstrings(String s, int k) {
    final int root = getRoot(k);
    int ans = 0;
    int vowels = 0;
    int vowelsMinusConsonants = 0;
    // {(vowels, vowelsMinusConsonants): count}
    Map<Pair<Integer, Integer>, Integer> prefixCount = new HashMap<>();
    prefixCount.put(new Pair<>(0, 0), 1);

    for (final char c : s.toCharArray()) {
      if (isVowel(c)) {
        vowels = (vowels + 1) % root;
        ++vowelsMinusConsonants;
      } else {
        --vowelsMinusConsonants;
      }
      Pair<Integer, Integer> prefix = new Pair<>(vowels, vowelsMinusConsonants);
      ans += prefixCount.getOrDefault(prefix, 0);
      prefixCount.merge(prefix, 1, Integer::sum);
    }

    return ans;
  }

  private boolean isVowel(char c) {
    return "aeiou".indexOf(c) != -1;
  }

  private int getRoot(int k) {
    for (int i = 1; i <= k; ++i)
      if (i * i % k == 0)
        return i;
    throw new IllegalArgumentException();
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class Solution:
  # Same as 2947. Count Beautiful Substrings I
  def beautifulSubstrings(self, s: str, k: int) -> int:
    kVowels = 'aeiou'
    root = self._getRoot(k)
    ans = 0
    vowels = 0
    vowelsMinusConsonants = 0
    # {(vowels, vowelsMinusConsonants): count}
    prefixCount = collections.Counter({(0, 0): 1})

    for c in s:
      if c in kVowels:
        vowelsMinusConsonants += 1
        vowels = (vowels + 1) % root
      else:
        vowelsMinusConsonants -= 1
      ans += prefixCount[(vowels, vowelsMinusConsonants)]
      prefixCount[(vowels, vowelsMinusConsonants)] += 1

    return ans

  def _getRoot(self, k: int) -> int:
    for i in range(1, k + 1):
      if i * i % k == 0:
        return i