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
37
38
39
class Solution {
 public:
  int beautifulPartitions(string s, int k, int minLength) {
    if (!isPrime(s.front()) || isPrime(s.back()))
      return 0;

    this->minLength = minLength;
    // dp[i][k] := # of beautiful partitions of s[i:] with k bars (|) left
    dp.resize(s.length(), vector<int>(k, -1));
    return partitions(s, minLength, k - 1);
  }

 private:
  static constexpr int kMod = 1'000'000'007;
  int minLength;
  vector<vector<int>> dp;

  int partitions(const string& s, int i, int k) {
    if (i <= s.length() && k == 0)
      return 1;
    if (i >= s.length())
      return 0;
    if (dp[i][k] != -1)
      return dp[i][k];

    // Not split between s[i - 1] and s[i].
    int ans = partitions(s, i + 1, k) % kMod;

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

    return dp[i][k] = ans % 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
34
35
36
37
class Solution {
  public int beautifulPartitions(String s, int k, int minLength) {
    if (!isPrime(s.charAt(0)) || isPrime(s.charAt(s.length() - 1)))
      return 0;

    this.minLength = minLength;
    // dp[i][k] := # of beautiful partitions of s.charAt(i:) with k bars (|) left
    dp = new Integer[s.length()][k];
    return partitions(s, minLength, k - 1);
  }

  private static final int kMod = 1_000_000_007;
  private int minLength;
  private Integer[][] dp;

  private int partitions(final String s, int i, int k) {
    if (i <= s.length() && k == 0)
      return 1;
    if (i >= s.length())
      return 0;
    if (dp[i][k] != null)
      return dp[i][k];

    // Not split between s.charAt(i - 1) and s.charAt(i).
    int ans = partitions(s, i + 1, k) % kMod;

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

    return dp[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
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:
      if i <= len(s) and k == 0:
        return 1
      if i >= len(s):
        return 0

      # Not 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)