Leetcode 0920 Number of Music Playlists

Problem statement

https://leetcode.com/problems/number-of-music-playlists/

Solution 1

Let $\phi(N, L, K)$ be answer for question: how many string of length $L$ we have, where we use $N$ symbols and distance between equal symbols is no more than $K$ (we can use less than all $N$ symbols. Then we can show, that $\phi(N, L, K) = N\cdot (N-1)\cdot\dots \cdot (N-K+1)\cdot (N-K)^{L-K}$. Next, we can use inclusion-exclusion formula to get answer: it is equal to \(\sum\limits_{i=1}^{N-K-1}\phi(N-i, L, K)\cdot C_N^i \cdot (-1)^i.\)

Complexity

Time complexity of this solution is O(LN).

Code

class Solution:
    def numMusicPlaylists(self, N, L, K):
        MOD = 10**9 + 7
        
        @lru_cache(None)
        def dfs_comb(i, j):
            if j == 0 or i == j: return 1
            return (dfs_comb(i-1, j) + dfs_comb(i-1, j-1)) % MOD
        
        def phi(N, L, K):
            ans = 1
            for i in range(N - K + 1, N + 1):
                ans = (ans*i) % MOD
            for i in range(L - K):
                ans = (ans*(N - K)) % MOD
            return ans
        
        return sum(phi(N-i, L, K) * dfs_comb(N, i) * (-1)**(i%2) for i in range(N-K)) % MOD

Remark

Actually, it can be done in O(N log L) complexity if we use fast exponentiation and modular arithmetic.

Solution 2

There is alternative dp solution: Let dp[i][j] be the number of play-lists of length i that have exactly j unique songs. We want dp[L][N], and it seems likely we can develop a recurrence for dp.

Consider dp[i][j]. Last song, we either played a song for the first time or we didn’t. If we did, then we had dp[i - 1][j - 1] * (N - j + 1) ways to choose it. If we didn’t, then we repeated a previous song in dp[i-1][j] * max(j-K, 0) ways (j of them, except the last K ones played are banned.)

Complexity

Time and space complexity is also O(NL).

Code

class Solution:
    def numMusicPlaylists(self, N, L, K):
        
        @lru_cache(None)
        def dp(i, j):
            if i == 0: return j == 0
            ans = dp(i-1, j-1) * (N-j+1)
            ans += dp(i-1, j) * max(j-K, 0)
            return ans % (10**9+7)

        return dp(L, N)