I've found an alternative solution to "AtCoder Grand Contest 032 B — Balanced Neighbors". For future use, here I write about the details about this problem. (URL to the problem: https://atcoder.jp/contests/agc031/tasks/agc031_b?lang=en)
I assume stones are placed in ascending order from left to right. For a stone at the position i (from now on I call this the i-th stone), let idr[i] denote the position of the stone which is right to the i-th stone and has the same color. For convenience, when no stones exist right to the i-th stone or idr[i] = i+1, I substitute -1 for idr[i] (-1 denotes "not defined"). For each i, idr[i] can be calculated using top-down (right-to-left) approach with O(N) time complexity. Then there will be two choices — paint between i-th and (idr[i])-th stones or not.
So I carried out simple top-down (right-to-left) dp.
- initial condition: dp(N) = 1
- dp(i) = dp(i+1) if idr[i] = -1, else dp(i+1) + dp(idr[i])
- answer: dp(0)
- time complexity: O(N)
- space complexity: O(N+M), where M denotes the number of colors (in this problem, M <= 2e5)