The first step is to find some criterion for which paintings can be created, where a painting is defined by $N$ numbers, each between $1$ and $M$ inclusive.

To this end, note that the last stamp Bessie uses will color $K$ consecutive units with the same color, and so in the final painting, there must be $K$ consecutive units with the same color.

Conversely, consider an arbitrary painting which satisfies this condition. It is not difficult to see that this painting must be attainable by some sequence of stampings: suppose the units in range $[i,i+K)$ have the same color. Start at the left end and work rightwards, stamping $[1,K+1)$ with the desired color for unit $1$, then stamping $[2, K+2)$ with the desired color for unit $2$, all the way until we reach $[i, K+i)$. Then similarly start from the right end and work leftwards. Once $[i, K+i)$ has been reached a second time, we have produced the desired painting.

So this problem is asking us to count the number of ways to pick $N$ numbers between $1$ and $M$ inclusive, so that some $K$ consecutive numbers are equal. As is often the case, it is simpler to count the complement. We will count the number of ways to pick $N$ numbers between $1$ and $M$ so that no $K$ consecutive numbers are all equal. Since there are $M^N$ ways to pick the numbers with no such restrictions, we will then subtract our complementary answer from $M^N$, to obtain our final answer.

We can use dynamic programming to solve this reduced problem. Let $\text{dp}(n)$ be the number of ways to pick $n$ numbers between $1$ and $M$ so that no $K$ consecutive numbers are equal. If $n<K$, this is a base case and the answer is $M^n$. Otherwise, note that in any good coloring, the last $K$ numbers cannot be equal. So for each good coloring, there is some $c < K$ so that the last $c$ numbers are equal, but the $c+1$-st number is different. Fix some $c$. Then there are $\text{dp}(n-c)$ ways to pick numbers for the first $n-c$ units, and $M-1$ ways to pick one number for the last $c$ units. This yields the recurrence relation

$$\text{dp}(n) = (M-1) \cdot \sum_{c=1}^{K-1} \text{dp}(n-c).$$

We immediately have a $O(NK)$ solution, which gets 75% of the points on this problem. To get full credit, one must make the following final observation. Let $s(n) = \sum_{i=1}^n \text{dp}(n)$. Then the above recurrence implies the following closed-form recurrence:

$$s(n) - s(n-1) = (M-1)(s(n-1) - s(n-K))$$

or
$$s(n) = Ms(n-1) - (M-1)s(n-K).$$

So rather than computing $\text{dp}(n)$ directly, we compute $s(n)$, and observe that $\text{dp}(N) = s(N) - s(N-1)$. This yields an $O(N)$ algorithm.

#include <iostream> #include <algorithm> using namespace std; #define MOD 1000000007 int s[10000001]; int main() { int N,M,K; cin >> N >> M >> K; s[0] = 0; for(int i=1;i<K;i++) s[i] = (M*((long long)s[i-1]) + M)%MOD; for(int i=K;i<=N;i++) s[i] = (M*((long long)s[i-1]) + MOD - ((M-1)*((long long)s[i-K]))%MOD)%MOD; int ans = 1; for(int i=1;i<=N;i++) ans = (M*((long long)ans))%MOD; cout << (((long long)ans) + MOD - ((long long)s[N]) + s[N-1])%MOD << '\n'; }