AOJ2730 Line Gimmick - ferinの競プロ帳

解法

ある矢印の並びでi番目の矢印からスタートするとする
i番目より左にある'>'とi番目より右にある'<'でのみ移動する方向を反転させる
sample1で考えてみる
2からスタートすると1で反転→4で反転→0で反転→5で反転→終了となる
3からスタートすると4で反転→1で反転→5で反転→0で反転→6で反転→終了となる

0123456
>><><<<

最後に反転する矢印が'<'であればここより右側の矢印が残る
同様に最後に反転する矢印が'>'であればここより左側の矢印が残る
各スタート位置について最後に反転する矢印の位置を高速で求められればうれしい
[0,i-1]の'>'の数と[i+1,n-1]の'<'の数がわかれば最終的な位置がわかりそう
これは累積和を使ってできる
'<'の数が'>'の数より多ければ最後に反転する矢印は'<'になる
これは右から数えて ('<'の数)-('>'の数) 番目の'<'になる
'>'の方が多い時も同様に考えられる
各スタート地点について最後に反転する矢印の位置がO(1)で求まるので全体O(n)で求まる

実装がめちゃくちゃつらかった…

ソースコード

'>'についての変数名にr、'<'についての変数名にlをつけています

#include <bits/stdc++.h>

using namespace std;
using ll = long long;
#define int ll
using VI = vector<int>;

#define FOR(i, a, n) for (ll i = (ll)a; i < (ll)n; ++i)
#define REP(i, n) FOR(i, 0, n)

template <typename T> T &chmin(T &a, const T &b) { return a = min(a, b); }
template <typename T> T &chmax(T &a, const T &b) { return a = max(a, b); }

signed main(void)
{
  cin.tie(0);
  ios::sync_with_stdio(false);

  int n;
  string s;
  cin >> n >> s;

  // sum_r[i] = ([0,i]の'>'の個数)
  // sum_l[i] = ([i,n-1]の'<'の個数)
  VI sum_r(n), sum_l(n);
  sum_l[n-1] = s[n-1]=='<'; sum_r[0] = s[0]=='>';
  FOR(i, 1, n) sum_r[i] = sum_r[i-1] + (s[i]=='>');
  for(int i=n-2; i>=0; --i) sum_l[i] = sum_l[i+1] + (s[i]=='<');

  // pos_r[i] = (左から数えてi番目の'>'の位置)
  // pos_l[i] = (右から数えてi番目の'<'の位置)
  VI pos_r(n), pos_l(n);
  int cnt = 0;
  REP(i, n) if(s[i]=='>') pos_r[cnt++] = i;
  cnt = 0;
  for(int i=n-1; i>=0; --i) if(s[i]=='<') pos_l[cnt++] = i;

  int ret = 1;
  REP(i, n) {
    // [i+1,n-1]の'<'の数、[0,i-1]の'>'の数
    int lnum = i==n-1?0:sum_l[i+1], rnum = i==0?0:sum_r[i-1];
    if(s[i]=='<') {
      if(lnum>=rnum) chmax(ret, pos_l[lnum-rnum]+1);  // 0-indexなので+1
      else chmax(ret, n-pos_r[rnum-lnum-1]);
    } else {
      if(rnum >= lnum) chmax(ret, n-pos_r[rnum-lnum]);
      else chmax(ret, pos_l[lnum-rnum-1]+1);  // 0-indexなので+1
    }
  }
  cout << ret << endl;

  return 0;
}