AtCoder Beginner Contest 172を見直す(その2）C - Tsundoku　尺取り法

問題

C - Tsundoku

机AとBにN冊とM冊の本がそれぞれ縦に積まれている、各々の本を読むための時間(Ai, Bi)が与えられるので、K分以内にできるだけ多く読んだ場合の冊数を答えよ。本は一番上にあるもののみ読むことができ、読み終わった本は縦積みには戻さないものとする。

制約

1<=N, M<=2*10^5
1<=K<=10^9
1<=Ai, Bi<=10^9

誤答

これは、一番上にある本を比較して早く読める方を読んでいけばええねや。と思って回答したのが以下。

M, N, K = map(int, input().split())
A = list(map(int, input().split()))
B = list(map(int, input().split()))
Ai, Bi, ans = 0, 0, 0
len_A, len_B = len(A), len(B)
MX = 1000000001
while K > 0:
   cst = 0
   cst_A = MX if Ai == len_A else A[Ai]
   cst_B = MX if Bi == len_B else B[Bi]
   if cst_A == cst_B == MX: break

   if cst_A < cst_B:
       cst = cst_A
       Ai += 1

当然「WA（Wrong Answer）」がでまくった。以下のような場合だと上手くいかない。

制限時間K：15分
Aの机の本：10, 1, 1, 1, 1, 1
Bの机の本：8, 8, 8

正しい出力は、「6」なのに、上記だと「2」が出力されてしまう。ということで考え直したんだけど、積まれた本の数が最大20万冊X2なので、全パターン試す場合は(2*10^5)^2となってどう考えても時間内に処理しきれそうにない。そしていい方法が思い付かずに終了となった。

動画解説の解答

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (int)(n); i++)
using namespace std;
using ll = long long;
using P = pair<int, int>;

int main() {
   int N, M, K;
   cin >> N >> M >> K;
   vector<int> A(N);
   vector<int> B(M);
   rep(i, N) cin >> A[i];
   rep(i, M) cin >> B[i];

   ll tot = 0;
   rep(i, M) tot += B[i];

   int j = M;
   int ans = 0;
   rep(i, N+1) {
       while (j > 0 && tot > K) {
           j -= 1;
           tot -= B[j];
       }
       if (tot > K) break;
       ans = max(ans, i+j);
       if (i == N) break;
       tot += A[i];
   }
   cout << ans << endl;

}

尺取り方というテクニックを使っているとのこと。ちなみにPythonだと以下のような感じ

N, M, K = map(int, input().split())
A = list(map(int, input().split()))
B = list(map(int, input().split()))
t, ans = sum(B), 0
j = M
for i in range(N+1):
   while j>0 and t>K:
       j -= 1
       t -= B[j]
   if t>K: break
   ans = max(ans, i+j)
   if i==N: break
   t += A[i]

print(ans)

写経はしてみたもののちゃんと理解できている気がしないので、尺取り法を使う問題を解いておこう。

解説動画

前回