鉄則本3-3 しゃくとり法 - かねぞうサバイブ

ついにしゃくとり法に来た。前々から小耳には挟んでいたけど実際に学ぶのは初めてなのでワクワクする。問題はこちら。

異なる2つの整数のペアを選んだ時、差がK以下であるものの総数を求める問題。全部でN(N-1)/2通りと書いてくれてるのはなかなか丁寧。N²オーダーでN=10⁵だと無理なので、全探索は無理だよと教えてくれている。しゃくとり法を使わなくても、今まで学んだ二分探索でいけそうなので、まずはそれで実装する。

回答コード

#include <bits/stdc++.h>
using namespace std;
#define dump(x) cout << #x << " = " << (x) << endl;
#define REP(i, n) for (ll i = 0; i < (n); i++)
#define OUT(n) cout << n << endl
typedef long long ll;
// #include <atcoder/all>
// using namespace atcoder;

int main() {
  ll N,K;cin>>N>>K;
  vector<ll> A(N);
  REP(i, N) cin>>A[i];
  ll ans = 0;
  REP(i, N) {
    ll h = A[i];
    ll L=0,R=N;
    while(R-L>1) {
      ll m = (L+R)/2;
      if (A[m]-h>K) {
        R=m;
      } else {
        L=m;
      }
    }
    ans += L-i;
  }
  OUT(ans);
}

配列の各要素について、その要素と比較する要素の差がKより上かどうかで二分していく。Kより上、つまり条件を満たさない要素は右に、条件を満たす要素は左に振り分けていく。そのため、最後にLが条件を満たす最大の要素の添字になっている。iを引いているのは既に探索した数を除くため。組み合わせのため、(Ai,Ai2)も(Ai2,Ai)も同じになる。

配列の各要素について二分探索しているので計算量はN(logN)、10⁵程度なら間に合う。

しゃくとり法

次はしゃくとり法を使って解く。初見なので本の解説を見ながら実装を進めていく。

#include <bits/stdc++.h>
using namespace std;
#define dump(x) cout << #x << " = " << (x) << endl;
#define REP(i, n) for (ll i = 0; i < (n); i++)
#define OUT(n) cout << n << endl
typedef long long ll;
// #include <atcoder/all>
// using namespace atcoder;

int main() {
  ll N,K;cin>>N>>K;
  vector<ll> A(N);
  REP(i, N) cin>>A[i];
  // R[i]
  vector<ll> R(N);
  ll ans = 0;
  for(ll i=0;i<N;i++) {
    if (i==0) R[i] = 0;
    else R[i] = R[i-1];
    while(A[R[i]+1]-A[i]<=K && R[i]+1<N) {
      R[i]++;
    }
    ans += R[i]-i;
  }
  OUT(ans);
}

At<=K+Ai となる最大のtをR[i]とする。このしゃくとり虫が条件を満たしていたら1ずつ進んでいくようにする。初期化の方法についてはiが0、つまり最初ならR[i]とし、それ以外ならR[i-1]としている。R[i-1]<R[i]なので、これは成り立つ。
計算量についてだが、R[i]を加算する処理はR[i]+1<Nより、全体でN-1回以下しか行われないため、定数倍となる。よって、外側のループのO(N)が計算量となる。つまり二分探索より計算量が少ない。すごい。