ABC071-C Together を解いてみた

今回のC問題はデータ構造を作ることができれば一瞬で解ける問題でした。

以下問題

問題

Together

実行時間制限: 2 sec / メモリ制限: 256 MB

問題文

長さ N の整数列 a1,a2,...,aN が与えられます。

各 1≤i≤Nに対し、ai に 1 足すか、1引くか、なにもしないかの三つの操作からどれか一つを選んで行います。

この操作の後、ある整数 Xを選んで、ai=X となる i の個数を数えます。

うまく操作を行い、Xを選ぶことで、この個数を最大化してください。

考えたこと

一つ一つ足して引いてのシミュレーションを行っていると O(N³) になってくるのでTLEになってしまいます。

視点を変えるとO(N)で問題が解けるようになります。

該当要素の＋1 - 1を + 1 して配列としてデータをもって置けばO(N)で計算が終了します。

この時配列に作るデータの注意点は +1 -1 が配列の要素外にならないようにする点です。

フロー

• 入力の受け取り
• 数列を受け取りながら以下の処理を行う
  • 数列の要素の値を受け取る
  • 要素 - 1 と 要素 + 1の変数を作る
  • -1 要素が0以上 かつ +1 要素が 10⁹ 以内ならば以下の処理
    • 3つの要素の該当部分を+1
  • -1要素が範囲外の場合
    • 該当要素と+1要素を+1
  • +1要素が範囲外の場合
    • 該当要素と-1要素を+1
  • 作った配列を線形的に探索
    • そのなかでの最大値をとる
• 最大値を出力

プログラム

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

int main(void){
  long long N;
  long long MAX = 100000;
  vector<long long> a(MAX);

  for(long long i = 0; i < MAX; i++){
    a[i] = 0;
  }
  
  cin >> N;

  for(long long i = 0;i < N; i++){
    long long tmp,dec,inc;
    cin >>  tmp;
    dec = tmp - 1;
    inc = tmp + 1;
    if(dec >= 0 && inc < MAX){
      a[dec]++;
      a[tmp]++;
      a[inc]++;
    }else if(dec < 0){
      a[tmp]++;
      a[inc]++;
    }else if(inc >= MAX){
      a[tmp]++;
      a[dec]++;
    }
  }
  long long mx = 0;
  for(long long i = 0; i < MAX; i++){
    mx = max(mx,a[i]);
  }
  cout << mx << endl;
  return 0;
}

こんな感じでACをとれることができました。