Skip to content

2478. Number of Beautiful Partitions 👍

  • Time: $O(nk)$
  • Space: $O(nk)$
 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
class Solution {
 public:
  int beautifulPartitions(string s, int k, int minLength) {
    if (!isPrime(s.front()) || isPrime(s.back()))
      return 0;
    vector<vector<int>> mem(s.length(), vector<int>(k, -1));
    return beautifulPartitions(s, minLength, k - 1, minLength, mem);
  }

 private:
  static constexpr int kMod = 1'000'000'007;

  // Returns the number of beautiful partitions of s[i..n) with k bars (|) left.
  int beautifulPartitions(const string& s, int i, int k, int minLength,
                          vector<vector<int>>& mem) {
    if (i <= s.length() && k == 0)
      return 1;
    if (i >= s.length())
      return 0;
    if (mem[i][k] != -1)
      return mem[i][k];

    // Don't split between s[i - 1] and s[i].
    int res = beautifulPartitions(s, i + 1, k, minLength, mem) % kMod;

    // Split between s[i - 1] and s[i].
    if (isPrime(s[i]) && !isPrime(s[i - 1]))
      res += beautifulPartitions(s, i + minLength, k - 1, minLength, mem);

    return mem[i][k] = res % kMod;
  }

  bool isPrime(char c) {
    return c == '2' || c == '3' || c == '5' || c == '7';
  }
};
 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
class Solution {
  public int beautifulPartitions(String s, int k, int minLength) {
    if (!isPrime(s.charAt(0)) || isPrime(s.charAt(s.length() - 1)))
      return 0;
    Integer[][] mem = new Integer[s.length()][k];
    return beautifulPartitions(s, minLength, k - 1, minLength, mem);
  }

  private static final int kMod = 1_000_000_007;

  // Returns the number of beautiful partitions of s[i..n) with k bars (|) left.
  private int beautifulPartitions(final String s, int i, int k, int minLength, Integer[][] mem) {
    if (i <= s.length() && k == 0)
      return 1;
    if (i >= s.length())
      return 0;
    if (mem[i][k] != null)
      return mem[i][k];

    // Don't split between s[i - 1] and s[i].
    int ans = beautifulPartitions(s, i + 1, k, minLength, mem) % kMod;

    // Split between s[i - 1] and s[i].
    if (isPrime(s.charAt(i)) && !isPrime(s.charAt(i - 1)))
      ans += beautifulPartitions(s, i + minLength, k - 1, minLength, mem);

    return mem[i][k] = ans % kMod;
  }

  private boolean isPrime(char c) {
    return c == '2' || c == '3' || c == '5' || c == '7';
  }
}
 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
class Solution:
  def beautifulPartitions(self, s: str, k: int, minLength: int) -> int:
    def isPrime(c: str) -> bool:
      return c in '2357'

    if not isPrime(s[0]) or isPrime(s[-1]):
      return 0

    kMod = 1_000_000_007

    @lru_cache(None)
    def dp(i: int, k: int) -> int:
      """
      Returns the number of beautiful partitions of s[i..n) with k bars (|)
      left.
      """
      if i <= len(s) and k == 0:
        return 1
      if i >= len(s):
        return 0

      # Don't split between s[i - 1] and s[i].
      ans = dp(i + 1, k) % kMod

      # Split between s[i - 1] and s[i].
      if isPrime(s[i]) and not isPrime(s[i - 1]):
        ans += dp(i + minLength, k - 1)

      return ans % kMod

    return dp(minLength, k - 1)