ABC 130 D問題

D - Enough Array

2019年6月16日開催のAtCoder Beginner Contest 130に参加した．D問題の解答例と解説．公式のPDF解説には尺取り法の実装例だけが載っていたので，このnoteでは二分探索も実装してみた．

D - Enough Array

問題文はこちら．与えられた長さNの正数数列Aの連続した部分列のうち，和が指定された値K 以上になるものが何通りか数える問題．Nが最大10^5なので，全探索では時間が足りないので一工夫必要．

数列Aの要素は全て正なので，ある部分列 a_left, a_(left+1), ... , a_right の和がK以上となる時，部分列，
a_left, a_left+1, ... , a_right
a_left, a_left+1, ... , a_right, a_right+1
a_left, a_left+1, ... , a_right, a_right+1, a_right+2
...
a_left, a_left+1, ... , a_right, a_right+1, a_right+2, ... a_N
も全て和がK以上になる．

そのため，leftを固定したとき，部分列a_left, a_(left+1), ... , a_rightの和がK以上となるような最小のrightを見つければ良い．このようなrightを探す方法として，尺取り法と二分探索がある．

尺取り法による解法

しゃくとりむしの動きのように部分列の範囲を動かして探索していく手法．left=1, right=1から始めて，部分列の和がK以上になるまでrightを1ずつ増やしていく．部分列の和がK以上になったら，その時のN-right+1をカウントし，leftを1増やす，を繰り返す．

C++による実装例は以下．(#include は省略)

#define rep(i,n) for(int i=0;i<n;i++)
using namespace std;

int main() {
 long long N, K;
 cin >> N >> K;
 vector<int> a(N,1);
 rep(i, N) {
   cin >> a[i];
 }
 
 // 尺取り法
 long long cnt = 0;
 int left = 0, right = 0; //和をとる部分列は[left, right] とする．
 long long sum = a[0]; // left=0, right=1 のとき
 
 while(right<N+1){ // rightがNを超えたら終了
   if(sum>=K) { //部分和がK以上
     //cout << left << ' ' << right << ' ' << sum << endl;
     cnt += N-right;  // right, right+1, ... , N-1 は全て条件を満たす
     
     left++; //leftを進める．
     sum -= a[left-1]; //進めた分，和を減らす
     
     if(right<left) {
       right++; //leftがrightを追い越した時，rightも進める
       sum += a[right]; //進めた分，和を増やす
     }
   } else { //部分和がKに足りない時は，rightを進める．
     right++;
     sum += a[right];//進めた分，和を増やす
   }
 }
 
 cout << cnt << endl;
 return 0;
}

二分探索による解法

探索範囲の中央で，条件より大きいか小さいかを判定し，次第に探索範囲を狭めていく手法．尺取り法よりも計算時間が長くなるが，公式pdf解説には実装例がなかったため実装してみた．

C++による実装例は以下．(#include は省略)​
部分和の計算では，累積和を用いて計算時間を短縮している．

#define rep(i,n) for(int i=0;i<n;i++)
using namespace std;

int main() {
 long long N, K;
 cin >> N >> K;
 vector<int> a(N,1);
 rep(i, N) {
   cin >> a[i];
 }
 
 // 累積和
 vector<long long> sum(N+1); // sum[i] = a[0] +... + a[i-1]
 sum[0] = 0;
 rep(i, N) {
   sum[i+1] = sum[i] + a[i];
 }
 
 // 二分探索
 long long cnt = 0;
 //和をとる部分列はsum[right]-sum[left] = a[left] +... + a[right-1]とする
 rep(left, N) {
   int start=left, end=N; // start, endは探索範囲の両端

   // 探索範囲内に和がK以上になるrightが存在するか
   if(sum[N]-sum[left]<K) break; // 存在しなければ打ち切り
 
   while(true) {
     int mid = (start+end)/2;
     if(sum[mid]-sum[left]>=K) {
       end = mid;
     } else {
       start = mid+1;
     }
     if(start==end) {
       //cout << left << ' ' << end << endl;
       cnt += N-end+1;
       break;
     }
   }
 }
 
 cout << cnt << endl;
 return 0;
}

計算時間比較

両方の実装例を提出してみた．実行時間最大ケースで，尺取り法は32ms，二分探索は35msとなり，今回のテストケースでは尺取り法の方が速いことが確認できた．

尺取り法

Submission #5999610 - AtCoder Beginner Contest 130

二分探索

Submission #5999631 - AtCoder Beginner Contest 130