Little Elephant and Retro Strings is an old problem. The intended solution is to use some intelligent grouping. I won't discuss it here because it's not relevant (and also to avoid spoiling), but you can read the details in the editorial. However, getting the right grouping idea is sometimes difficult, especially in problems with multiple parameters. Hence, I decided to go with a naive DP:

`dp[progress][color][length]`

is the number of strings with a suffix of given `color`

and `length`

. Here `progress`

is 0 if we haven't seen `"B...B"`

before, 1 if we have seen `"B...B"`

but not `"W...W"`

, and 2 if we have seen both. The answer is the sum of `dp[2]`

. Notice that our position in the string is not a part of the state because we process the updates in place.

How do we process updates in such a dp? When appending a `'B'`

, all strings ending with `'W'`

terminate, and all strings ending with `'B'`

get one longer. Since amortized analysis applies here, the strings that terminate can be processed with a naive loop. However, strings that get one longer seemingly require moving all values in our dp by one position. Fortunately, we can do this in O(1) if we use a queue (or a deque) to store our dp values. Instead of moving all values by one, we can just prepend a 0 to our dp. Appending a `'W'`

follows the same process.

However, when considering an `'X'`

, we cannot use amortized analysis because we cannot erase values. This is because every string has a valid continuation: strings ending with `'W'`

can get another `'W'`

, and strings ending with `'B'`

can get another `'B'`

. Therefore, we need to speed up the updates where the last color changes. Specifically, such updates require us to increment `dp[progress][!color][1]`

by the sum of `dp[progress][color]`

. Therefore, it is sufficient to introduce another dp: `sum_dp[progress][color] = sum(dp[progress][color])`

. There is only a constant number of transitions, so maintaining `sum_dp`

is not an issue either.

You can examine my submission for details. I may clean it up if I get some free time. I've been inspired by this excellent video. The takeaway is that using alternative data structures for your DP can speed up certain operations. Shifting the DP by one is just one example of what one may achieve with it.