Leetcode: 2478. Number of Beautiful Partitions

• Pattern: Counting Problem.
• Pattern: Find optimal cuts with length restrictions.

• Can we state the problem in terms of number of X starting on any index greater than Y?

Be \(dp(i, k)\) the number of beautiful partitions where the first partition starts on any index between \(i\) and \(n-1\). Base cases: if \(i>n\) or \(k<0\), the answer is 0. If \(i=n\), the answer is 1 if \(k=0\) and 0 otherwise. Induction step: if \(s[i]\) is a prime and \(s[i-1]\) is not a prime, then \(dp(i, k)=dp(i+1,k) + dp(i + minLength, k - 1)\) since we must consider cutting on \(i\). Otherwise, \(dp(i, k)=dp(i + 1, k)\) since there is no cut to consider at position \(i\). Time and space complexity is \(O(n^2)\).

class Solution:
    def beautifulPartitions(self, s: str, k: int, minLength: int) -> int:
        N = len(s)
        P = "2357"

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

        @cache
        def dfs(i, k, level=0):
            if i > N or k < 0:
                return 0
            if i == N:
                return 1 if k == 0 else 0
            ans = dfs(i + 1, k, level + 1)
            if s[i - 1] not in P and s[i] in P:
                ans += dfs(i + minLength, k - 1, level + 1)
            return ans % (10**9 + 7)

        return dfs(minLength, k - 1)

• Can we derive an invariant based on the smallest possible examples? If \(s\) starts is a non-prime or ends with a prime, then the solution is always 0.
• Can we pre-process the input in a way to make easy to solve the problem? We want to partition \(s\) only in positions where \(s[i]\) is a non-prime and \(s[i+1]\) is a prime. So, we can create \(starts\) with indexes where valid partitions can start. Now, the problem becomes find partitions on indexes in \(starts\) of length at least \(minLength\).

• The first idea that comes to mind is using Dynamic Programming to solve the problem. Be \(dp(i, k)\) the number of ways to partitions \(s[i...(n-1)]\) in \(k\) partitions. Be \(V(i)\) indexes \(j\) where \(s[i...(j-1)]\) is a valid partition. It is not difficult to see how to compute each sub problem in linear time. This solution won't make in time, since the complexity is \(O(n^3)\).

  \[ dp(i, k)=\begin{cases} 0, & \mbox{if $i = n$ and $k \neq 0$} \\ 1, & \mbox{if $i = n$ and $k = 0$} \\ \sum_{j \in V(i)} dp(j, k - 1), & \mbox{otherwise}. \end{cases} \]

• Can we break-down the problem in small and easily to solve parts? The problem with the above recurrence is that it takes linear time to compute each sub problem. To improve it, we can break-down the third-case in other Dynamic Programming that depends on \(dp\). Be \(ac(i, k)\) the number of beautiful partitions starting on $starts[i], starts[i+1], starts[i+2], …$. We have that \(ac(i, k)=\sum_{i \leq j < |starts|} dp(i, k)\). Note that there are \(O(n^2)\) possible states for \(ac\), but we can compute each one in \(O(1)\). Now, we can re-write \(dp\) in terms of \(ac\) and the time complexity to solve the problem becomes \(O(n^2 \log n)\) with space \(O(n^2)\). The log comes from the Binary Search to find the next valid position to start a partition after \(i\).

  \[ dp(i, k)=\begin{cases} 0, & \mbox{if $i = n$ and $k \neq 0$} \\ 1, & \mbox{if $i = n$ and $k = 0$} \\ ac(j, k - 1), & \mbox{where $j$ is the first valid start of partition after $i$}. \end{cases} \]

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

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

        starts = []
        for i, a, b in zip(range(N), s, s[1:]):
            if a not in primes and b in primes:
                starts.append(i + 1)
        starts.append(N)

        @cache
        def dfs2(i, k):
            if i == len(starts):
                return 0
            return (dfs1(starts[i], k) + dfs2(i + 1, k)) % MOD

        @cache
        def dfs1(i, k):
            if i == N:
                return 1 if k == 0 else 0
            if k < 0:
                return 0
            return dfs2(bisect_left(starts, i + minLength), k - 1) % MOD

        return dfs1(0, k)

• Interview preparation: V