LeetCode 2478 - Number of Beautiful Partitions

Difficulty: 🔴 Hard
Topics: String, Dynamic Programming, Prefix Sum

Solution

Problem Understanding

The problem asks us to count the number of ways to split a given string s of digits into exactly k non-overlapping substrings, where each substring satisfies specific rules: it must start with a prime digit (2, 3, 5, 7), end with a non-prime digit (1, 4, 6, 8, 9), and have length at least minLength. The result must be returned modulo 10^9 + 7 because the number of partitions can be very large.

The input consists of a string s of digits from '1' to '9', an integer k for the number of required partitions, and minLength, the minimum allowed length for each substring. The constraints guarantee that s is relatively small (length <= 1000), so algorithms up to roughly O(n^2) or O(n^2 * k) may be feasible, but brute-force enumeration of all partitions is likely too slow.

Edge cases include situations where the string is too short to produce k substrings of at least minLength, or the first or last digits do not satisfy the prime/non-prime rules, which immediately makes the answer zero.

Approaches

Brute Force

The naive approach would try every possible way to split the string into k substrings and check each candidate partition for the prime-start, non-prime-end, and length constraints. While this guarantees correctness, it is extremely inefficient because the number of ways to partition a string of length n into k substrings is combinatorial (C(n-1, k-1)), which can be very large for n = 1000.

Optimal Approach

The key observation is that the constraints can be checked incrementally and the problem has optimal substructure: if we know the number of beautiful partitions ending at a certain index with j partitions, we can build larger partitions by extending from valid previous indices. This suggests a dynamic programming solution.

We define dp[i][j] as the number of beautiful partitions that end at index i of s with exactly j substrings. For each valid start index, we check whether a substring of length at least minLength ending at i is valid. Using prefix sums of the DP array allows us to efficiently compute transitions in O(1) per index.

This reduces the brute-force combinatorial search to a manageable dynamic programming solution with time complexity around O(n * k).

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(C(n-1, k-1) * n)	O(k)	Enumerates all partitions and checks constraints
Optimal	O(n * k)	O(n * k)	Uses dynamic programming with prefix sums to track valid partitions

Algorithm Walkthrough

Identify prime digits as a set for constant-time lookup: {'2', '3', '5', '7'}.
Immediately check if the first digit of s is prime and the last digit is non-prime. If not, return 0 because no valid partition is possible.
Initialize a DP table dp[i][j] where i is the end index in s and j is the number of partitions.
Set dp[i][1] to 1 for all valid single substrings starting at index 0 with length at least minLength and satisfying the start/end rules.
For each index i from 0 to len(s)-1 and for each partition count j from 2 to k, iterate backward to find valid previous partition endpoints that can form a new substring ending at i of length at least minLength.
Use prefix sums to sum all valid dp[prev][j-1] efficiently instead of looping.
Return the total number of partitions ending at len(s)-1 with exactly k substrings modulo 10^9 + 7.

Why it works: The dynamic programming invariant is that dp[i][j] correctly counts all partitions ending at i with j substrings. Each transition only considers valid previous substrings that satisfy the constraints, ensuring correctness.

Python Solution

class Solution:
    def beautifulPartitions(self, s: str, k: int, minLength: int) -> int:
        MOD = 10**9 + 7
        n = len(s)
        primes = {'2', '3', '5', '7'}

        if s[0] not in primes or s[-1] in primes:
            return 0
        
        valid = [0] * n
        for i in range(minLength - 1, n):
            if s[i] not in primes and s[i - minLength + 1] in primes:
                valid[i] = 1

        dp = [[0] * (k + 1) for _ in range(n)]
        for i in range(minLength - 1, n):
            if valid[i]:
                dp[i][1] = 1

        for j in range(2, k + 1):
            prefix = [0] * n
            prefix[0] = dp[0][j - 1]
            for i in range(1, n):
                prefix[i] = (prefix[i - 1] + dp[i][j - 1]) % MOD

            for i in range(minLength - 1, n):
                if not valid[i]:
                    continue
                left = i - minLength
                if left < 0:
                    continue
                dp[i][j] = prefix[left]

        return sum(dp[i][k] for i in range(n)) % MOD

Implementation Explanation: We first mark all valid endpoints of substrings in valid. The DP table dp[i][j] counts partitions ending at i with j parts. Prefix sums prefix efficiently accumulate sums of previous DP values to avoid nested loops. Finally, we sum all valid endpoints with k partitions.

Go Solution

func beautifulPartitions(s string, k int, minLength int) int {
    const MOD = 1_000_000_007
    n := len(s)
    primes := map[byte]bool{'2': true, '3': true, '5': true, '7': true}

    if !primes[s[0]] || primes[s[n-1]] {
        return 0
    }

    valid := make([]int, n)
    for i := minLength - 1; i < n; i++ {
        if !primes[s[i]] && primes[s[i-minLength+1]] {
            valid[i] = 1
        }
    }

    dp := make([][]int, n)
    for i := range dp {
        dp[i] = make([]int, k+1)
    }

    for i := minLength - 1; i < n; i++ {
        if valid[i] == 1 {
            dp[i][1] = 1
        }
    }

    for j := 2; j <= k; j++ {
        prefix := make([]int, n)
        prefix[0] = dp[0][j-1]
        for i := 1; i < n; i++ {
            prefix[i] = (prefix[i-1] + dp[i][j-1]) % MOD
        }
        for i := minLength - 1; i < n; i++ {
            if valid[i] == 0 {
                continue
            }
            left := i - minLength
            if left < 0 {
                continue
            }
            dp[i][j] = prefix[left]
        }
    }

    ans := 0
    for i := 0; i < n; i++ {
        ans = (ans + dp[i][k]) % MOD
    }
    return ans
}

Go Notes: Go implementation mirrors Python logic, with slices and explicit initialization. Map is used for constant-time prime lookup. Integer overflows are handled by % MOD during accumulation.

Worked Examples

Example 1: `s = "23542185131", k = 3, minLength = 2`

i (end index)	valid[i]	dp[i][1]	dp[i][2]	dp[i][3]
3	1	1	0	0
5	1	1	1	0
8	1	1	2	1
10	1	1	2	3

Total partitions with 3 parts = 3

Example 2: `s = "23542185131", k = 3, minLength = 3`

Total partitions with 3 parts = 1

Example 3: `s = "3312958", k = 3, minLength = 1`