Editorial for Mock CCC '19 Contest 1 J5 - Pusheen Designs a Sushi Boat
Submitting an official solution before solving the problem yourself is a bannable offence.
For the first subtask, the answer is always ~N~.
For the second subtask, there is an ~\mathcal O(K^2)~ solution that consists of looping over all pairs of pieces of sashimi and counts the number of pairs which come from distinct fish.
To get full credit, the above solution cannot be directly generalized as that would run in ~\mathcal O(N^K)~. We instead need to be more careful in how we enumerate these options.
Imagine that we loop over the fish in order from ~1~ to ~N~, and at each point in time we decide whether to select one piece of sashimi from that fish or not. If we take this approach, then the only meaningful component of our state is how many pieces of sashimi we have taken.
Define ~f(n, k)~ to be the number of ways to select ~k~ pieces of sashimi using only pieces from the first ~n~ fish. We have that ~f(0, 0) = 1~ and ~f(n, k) = f(n-1, k) + a[n] \cdot f(n-1, k-1)~ where we define ~a[n]~ to be the number of pieces of sashimi from fish ~n~ that we can select.
With memoization, this DP takes ~\mathcal O(NK)~ space and runs in ~\mathcal O(NK)~ time.