E1. Close Tuples (easy version)

• Quiz
  • https://codeforces.com/contest/1462/problem/E1
• AC
  • https://codeforces.com/contest/1462/submission/101418838
• 解説
  • editorialと違う解法だったので書いておく
  • 3つ組なので真ん中をずらしながら全探索する
  • 左側の統計と右の統計をBITでそれぞれ持てばいい。下記画像に合わせて具体例を言うと、真ん中が4なら、左が取れる値は2～6, そのそれぞれについて、右側が取れる範囲も決まるので数えられる

void solve(){
  ll N;
  cin>>N;
  VI A(N);
  rep(i, N){
    cin >> A[i];
  }

  BIT left(N+10);
  BIT right(N+10);

  left.add(A[0],1);
  FOR(i,2,N){
    ll a = A[i];
    right.add(a,1);
  }
  ll sum=0;
  FOR(i,1,N-1){
    // iが真ん中
    ll v = A[i];

    sum += left.range_sum(v-2,v-2) * right.range_sum(v-2,v);
    sum += left.range_sum(v-1,v-1) * right.range_sum(v-2,v+1);
    sum += left.range_sum(v  ,v  ) * right.range_sum(v-2,v+2);
    sum += left.range_sum(v+1,v+1) * right.range_sum(v-1,v+2);
    sum += left.range_sum(v+2,v+2) * right.range_sum(v  ,v+2);

    left.add(v,1);
    right.add(A[i+1],-1);
  }
  p(sum);
}

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

    // input
    ll N; 
    cin>>N;
    
    while(N--)solve();
    
    return 0;
}