leetcode - 920. Number of Music Playlists

Description

Your music player contains n different songs. You want to listen to goal songs (not necessarily different) during your trip. To avoid boredom, you will create a playlist so that:

Every song is played at least once.
A song can only be played again only if k other songs have been played.
Given n, goal, and k, return the number of possible playlists that you can create. Since the answer can be very large, return it modulo 10^9 + 7.

Example 1:

Input: n = 3, goal = 3, k = 1
Output: 6
Explanation: There are 6 possible playlists: [1, 2, 3], [1, 3, 2], [2, 1, 3], [2, 3, 1], [3, 1, 2], and [3, 2, 1].

Example 2:

Input: n = 2, goal = 3, k = 0
Output: 6
Explanation: There are 6 possible playlists: [1, 1, 2], [1, 2, 1], [2, 1, 1], [2, 2, 1], [2, 1, 2], and [1, 2, 2].

Example 3:

Input: n = 2, goal = 3, k = 1
Output: 2
Explanation: There are 2 possible playlists: [1, 2, 1] and [2, 1, 2].

Constraints:

0 <= k < n <= goal <= 100

Solution

Ref: https://leetcode.com/problems/number-of-music-playlists/solutions/3869515/c-simple-solution-dp-memoization-with-explaination/

For the current song, it’s either pick or not pick. Let dp[i][j] denotes the solutions when we have i different songs, and the length of the list is j.
If we pick the ith song, then we have i choices, the whole count will be i * dp[i - 1][j - 1]
If we don’t pick the ith song, we could pick from the already picked songs, then we have i - k choices, the whole count will be (i - k) * dp[i][j - 1]
So the transformation equation is:
dp[i][j] = i * dp[i - 1][j - 1] + (i - k) * dp[i][j - 1]

Time complexity: $o(N\ast L)$
Space complexity: $o(N\ast L)$

Code

class Solution:
    def numMusicPlaylists(self, n: int, goal: int, k: int) -> int:
        dp = [[0] * (goal + 1) for _ in range(n + 1)]
        dp[0][0] = 1
        modulo = 1000000007
        for i in range(1, n + 1):
            for j in range(1, goal + 1):
                dp[i][j] = (i * dp[i - 1][j - 1] + max(i - k, 0) * dp[i][j - 1]) % modulo
        return dp[-1][-1]