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)

My solution

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)

Source code: https://atcoder.jp/contests/agc031/submissions/7460148?lang=en