Problem statement

https://leetcode.com/problems/count-ways-to-distribute-candies/

Solution

This problem is about ealuating stirling numbers of the second kind. We have reccurent formula dp(n, k) = dp(n-1, k)*k + dp(n-1, k-1): when we add element with index n we can either put it in separate set with dp(n-1, k-1) number of ways or we add it to one of the k sets with dp(n-1, k)*k options. Also we have border cases: dp(0, 0) = 1 and dp(0, k) = dp(n, 0) = 0 with n != 0 and k != 0.

Complexity

Time complexity is O(nk), for some reason it is TLE, so we need to reduce cache.

Code

class Solution:
    def waysToDistribute(self, n, k):
        @lru_cache(1000)
        def dp(n, k):
            if n == 0 and k == 0: return 1
            if n <= 0 or k <= 0: return 0
            return (dp(n-1, k)*k + dp(n-1, k-1)) % (10**9 + 7)

        return dp(n, k)

Remark

There is a way to compute number in O(k log n) complexity, using the following formula \(S(n, k) =\sum _{j=1}^{k}(-1)^{k-j}{\frac {j^{n-1}}{(j-1)!(k-j)!}}\)