yokobuttonの不定期で競技プログラミングをするブログ

不定期で解けた競技プログラミングコンテストの問題を載せています。

トヨタシステムズプログラミングコンテスト2021(AtCoder Beginner Contest 228) D - Linear Probing

問題の要約
 N=2^20項からなる数列A=(A0,A1,...,AN-1)があり,はじめ,全ての要素は-1。
 Q個のクエリを順番に処理せよ。
 i(1<=i<=Q)個目のクエリはti=1またはti=2を満たす整数ti及び整数xiで表され,内容は以下の通り。
  ・ti=1のとき,以下の処理を順番に行う。
   1,整数hをh=xiで定める。
   2,AhmodN≠-1である間,hの値を1増やすことを繰り返す。この問題の制約下でこの操作が有限回で終了することは証明できる。
   3,AhmodNの値をxiで書き換える。
  ・ti=2のとき,その時点でのAximodNの値を出力する。
制約
 1<=Q<=2*10^5
 ti∈{1,2}(1<=i<=Q)
 0<=xi<=10^18(1<=i<=Q)
 ti=2であるようなi(1<=i<=Q)が1つ以上存在する。 
入力
 Q
 t1 x1
 .
 .
 .
 tQ xQ
考え方
 1,ti=2のときはA[x%N]をそのまま出力すればよい。
 2,ti=1のときを考える。
  vector<long long> next(N)を用意して,次に探す番号を入れればよさそう。next[i]の初期値はi。
  nextのx%Nの次を探すfind関数を用意するとよさそう。
  find関数はnext[x%N]=x%Nならばx%Nを返す。それ以外の時はnext[x%N]の次を探してあげる。同時に返り値をnext[x%N]に代入する。
  find関数で見つけた番号をiとし,A[i]にxを代入する。
  next[i]に,find関数の引数に(i+1)%Nを入れた返り値を代入する。
 

実際のプログラム
#include<iostream>
#include<vector>
#include<map>

using namespace std;

const long long N = 1048576;

long long find(vector<long long> &next, long long x){
  if(next[x] == x){
    return x;
  }else{
    return next[x] = find(next, next[x]);
  }
}

int main(){
  long long Q;
  cin >> Q;
  vector<long long> A(N, -1);
  vector<long long> next(N);
  for(long long i = 0; i < N; i++){
    next[i] = i;
  }
  for(long long q = 0; q < Q; q++){
    long long t, x;
    cin >> t >> x;
    if(t == 1){
      long long i = find(next, x%N);
      A[i] = x;
      next[i] = find(next, (i+1)%N);
    }else if(t == 2){
      cout << A[x%N] << endl;
    }
  }
  
  return 0;
}