[Q] B - Arithmetic Progression Subsequence

考察
0. A<=10と特徴的。何かAに絞って数え上げられないか考える。
O(NlogN*10)くらいでいけそうだと思う。

1. index_iを0~N-1まで見るとして。数列の末端をindex_iに固定して考える。
2. 等差数列のパターンを出してみると、どの数も4~5通りしかないことに気づく。

  // 等差数列の組み合わせを調べる
  // 1だとして、moves[1] = {111,321,531,741,951}
  // 2だとして、moves[2] = {222,432,642,852}
  // 3だとして、moves[3] = {123,333,543,753,963}

3. なので、A[i]を等差数列の末端要素として使うと仮定して、等差数列ができるかどうかを調べればいい。

4. 調べ方は、あらかじめ数字の出現indexを調べておけばできそう。
A[i] = 5で{1,3,5}を調べたい場合、
a. 3がiよりも手前のindexで出現しているか調べ、index_3を見つけて、
b. 1がindex_3よりも手前のindexで出現しているか調べ、index_1を見つける。
[index_1, i]が、等差数列を含む部分列の最小幅なので、[0, index_1]の開始地点の選び方が答えのバリエーションになる。

4. 等差数列ができたときにメモすべき始点のindexは、できるだけ後ろにあるものを選ぶのが最善。

入力例1
5
5 3 4 1 5

1. A[3] = 1
5 3 4 1
~ ~   ~
5 3   1 ができることがわかる。
5の出現位置はindex_0なので、[0, 3] の1通りが答えに加算。

2. A[4] = 5
5 3 4 1 5
  ~ ~   ~
  3 4   5 ができることがわかる。
3の出現位置はindex_1なので、[1, 4],[0, 4]の2通りが加算。

答えは3通り。


サンプル
6
1 1 2 2 3 10

1. A[4] = 3を見たとき。
1 1 2 2 3
  ~   ~ ~
1,2,3をつくることができる。これは、
A[1]=1を選ぶことで、[1, 4]と[0, 4]の2通りが選べると確定できる。
A[0]=1は選ばない。

2. A[5] = 10を見るとき。
1 1 2 2 3 10
  ~   ~ ~
10で等差数列は作れないけど、A[1] = 1を始点にすれば{1,2,3}が作れることが分かっているので、
[1, 5], [0, 5]の2通りを加算。

答えは4通り

実装

int main() {
  cincout();
  
  ll N;
  cin >> N;
  vector<ll> A(N);
  rep(i, N) cin >> A[i];

  // 数字の出現indexを管理
  vector<ll> G[11];
  rep(i, N){
    G[A[i]].push_back(i);
  }

  // 等差数列の組み合わせを調べる
  // 1だとして、moves[1] = {111,321,531,741,951}
  // 2だとして、moves[2] = {222,432,642,852}
  // 3だとして、moves[3] = {123,333,543,753,963}
  vector<tll> moves[11];
  for(ll i = 1; i <= 10; ++i) for(ll j = 1; j <= 10; ++j) for(ll k = 1; k <= 10; ++k){
    if (j - i == k - j){
      moves[k].push_back(tll(i, j, k));
    }
  }

  ll ans = 0;
  ll max_start = -1;

  rep(i, N){ // 選ぶ数列の末端をiに固定する。
    ll aa = A[i];
    for(auto[a, b, c]: moves[aa]){
      // iより少し前にあるbのindex_bを探して
      // index_bより少し前にあるaのindex_aを探す。
      
      auto bit = lower_bound(all(G[b]), i);
      if (bit == G[b].begin()) continue;
      --bit;
      ll ib = *bit;

      auto ait = lower_bound(all(G[a]), ib);
      if (ait == G[a].begin()) continue;
      --ait;
      ll ia = *ait;

      // 等差数列が見つかったので、できるだけ後ろの開始地点をメモ
      chmax(max_start, ia);
    }
    if (max_start > -1){
      // [max_start, i]を選べば、等差数列を含むことが分かった。
      // 答えになるバリエーションは、開始地点[0, max_start]通りと、末端にiを選んだときの組み合わせ。
      ans += max_start + 1;
    }
  }
  cout << ans << endl;
}

https://atcoder.jp/contests/arc170/submissions/49550725