2478. Number of Beautiful Partitions

Description

You are given a string s that consists of the digits '1' to '9' and two integers k and minLength.

A partition of s is called beautiful if:

• s is partitioned into k non-intersecting substrings.
• Each substring has a length of at least minLength.
• Each substring starts with a prime digit and ends with a non-prime digit. Prime digits are '2', '3', '5', and '7', and the rest of the digits are non-prime.

Return the number of beautiful partitions of s. Since the answer may be very large, return it modulo 109 + 7.

A substring is a contiguous sequence of characters within a string.

 

Example 1:

Input: s = "23542185131", k = 3, minLength = 2
Output: 3
Explanation: There exists three ways to create a beautiful partition:
"2354 | 218 | 5131"
"2354 | 21851 | 31"
"2354218 | 51 | 31"

Example 2:

Input: s = "23542185131", k = 3, minLength = 3
Output: 1
Explanation: There exists one way to create a beautiful partition: "2354 | 218 | 5131".

Example 3:

Input: s = "3312958", k = 3, minLength = 1
Output: 1
Explanation: There exists one way to create a beautiful partition: "331 | 29 | 58".

 

Constraints:

• 1 <= k, minLength <= s.length <= 1000
• s consists of the digits '1' to '9'.

Solution

number-of-beautiful-partitions.py

class Solution:
    def beautifulPartitions(self, s: str, k: int, minLength: int) -> int:
        N = len(s)
        M = 10 ** 9 + 7

        primes = set(["2", "3", "5", "7"])

        if s[0] not in primes or s[-1] in primes: return 0

        @cache
        def go(index, parts):
            if parts == 0 and index <= N: return 1
            if index >= N: return 0

            res = go(index + 1, parts)

            if s[index] in primes and s[index - 1] not in primes:
                res += go(index + minLength, parts - 1)

            return res % M

        return go(minLength, k - 1)

number-of-beautiful-partitions.cpp

class Solution {
public:
    const int M = 1e9 + 7;
    int cache[1001][1001];

    bool isPrime(char x) {
        return x == '2' || x == '3' || x == '5' || x == '7';
    }

    int go(int index, int parts, string& s, int k, int minLength, int N) {
        if (cache[index][parts] != -1) return cache[index][parts];

        if (parts >= k) return cache[index][parts] = 0;

        if (parts == k - 1) return cache[index][parts] = (N - index >= minLength);

        int res = 0;

        for (int j = index + minLength; j < N - (k - parts - 1) * minLength + 1; j++) {
            if (isPrime(s[j]) && !isPrime(s[j - 1])) {
                res += go(j, parts + 1, s, k, minLength, N);
                res %= M;
            }
        }

        return cache[index][parts] = res;
    }

    int beautifulPartitions(string s, int k, int minLength) {
        int N = s.size();
        memset(cache, -1, sizeof(cache));

        if (!isPrime(s[0]) || isPrime(s[N - 1])) return 0;

        return go(0, 0, s, k, minLength, N);
    }
};