Codeforces Round #591 D. Stack Exterminable Arrays

問題

https://codeforces.com/contest/1240/problem/D

解法

解説を見た。

各indexについてそこから順にstackに入れていった時に初めて空になるindex nxt[i]が欲しくなる(これが分かればあとはdpで数えるだけ)。 しかし直接nxtを求めるには少し情報が足りないので補助的に次の配列を考える

nxtX[i][x] := a[i]...a[j]が空になって、かつa[j+1] = xであるような最小のj>i

これを必要なところだけHashMapで持って更新していけばOK。nxtXの更新はコピーが発生しないようにswapなりmoveを使うことに注意すればO(n)でnxt,nxtXを求められる。

a[0]...a[i]をスタックに入れた状態を文字列と見てハッシュにして同じ状態のペアを数える方法もあるみたい。こっちのほうが思考停止で書けるかも。

参考: Codeforces Round #591 Div2 F / Div1 D. Stack Exterminable Arrays • knshnbのブログ

void solve() {
  int n;
  cin >> n;
  vector<int> a(n);
  cin >> a;
  vector<int> nxt(n + 1, -1);
  vector<unordered_map<int, int>> nxtX(n + 1);
  vector<ll> cnt(n + 1, 0ll);
  ll ans = 0ll;
  for (int i = n - 1; i >= 0; i--) {
    if (i + 1 < n && a[i] == a[i + 1])
      nxt[i] = i + 1;
    else if (nxtX[i + 1].find(a[i]) != nxtX[i + 1].end())
      nxt[i] = nxtX[i + 1][a[i]] + 1;

    if (nxt[i] != -1) swap(nxtX[i], nxtX[nxt[i] + 1]);
    if (nxt[i] != -1 && nxt[i] + 1 < n) nxtX[i][a[nxt[i] + 1]] = nxt[i];
    if (nxt[i] != -1) cnt[i] = 1 + cnt[nxt[i] + 1];
    ans += cnt[i];
  }
  cout << ans << endl;
  return;
}