【AtCoder】ABC 430 C - Truck Driver

C - Truck Driver

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

AtCoder

実行時間制限: 2 sec / メモリ制限: 1024 MiB / Difficulty: 750 / NoviSteps: 2Q / 配点: 300 点

問題概要

a, bからなる長さ $N$ の文字列 $S$ と正整数 $A,B$ が与えられる。 以下の条件を全て満たす整数組 $(l,r)$ の個数を求めよ。

• $1\leq l \leq r \leq N$
• $S$ の $l$ 文字目から $r$ 文字目までに含まれるaの個数が $A$ 以上
• $S$ の $l$ 文字目から $r$ 文字目までに含まれるbの個数が $B$ 未満

制約

• $1\leq N \leq 3\times 10^5$
• $1 \leq A,B \leq N$
• 与えられる数値は全て整数

考察

$l$ を固定して、条件を満たす $r$ の個数を数え上げる典型問題である。

このときに、任意の区間内のaとbの個数を何度も計算する必要があるので、累積和を用いて高速に求められるようにしておく。

• $a_i :=$ $S$ の先頭から $i$ 文字目までに含まれるaの個数
• $b_i :=$ $S$ の先頭から $i$ 文字目までに含まれるbの個数

とすれば、$S$ の $l$ 文字目から $r$ 文字目までに含まれるaの個数は $a_r - a_{l-1}$、bの個数は $b_r - b_{l-1}$ として、計算量 $O(1)$ で求めることができる。

ここで $l \: (1 \leq l \leq N)$ をある値に固定したとき、 $r$ が満たすべき条件は以下のようになる。

• $a_{r_a} - a_{l-1} \geq A \iff a_{r_a} \geq a_{l-1} + A$
• $b_{r_b} - b_{l-1} < B \iff b_{r_b} < b_{l-1} + B$

累積和テーブルは単調増加なので、上記の条件を満たす $r_a, r_b$ はそれぞれ二分探索で求めることができる。

$r_a \leq r_b$ であれば、この $l$ に対して条件を満たす $r$ の個数は $r_b - r_a + 1$ 個となるので、これを全ての $l$ について合計すればよい。

全体の計算量は $O(N \log N)$ である。

実装例

イテレータとインデックスの変換、そして $r_a, r_b$ の求め方に注意すること。

1.#include <bits/stdc++.h>
2.using namespace std;
3. 
4.#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)
5. 
6.using ll = long long;
7. 
8.// ======================================== //
9. 
10.int main()
11.{
12.    int N, A, B;
13.    string S;
14.    cin >> N >> A >> B >> S;
15. 
16.    vector<int> ps_a(N + 1, 0), ps_b(N + 1, 0);
17.    rep(i, 0, N)
18.    {
19.        if (S[i] == 'a')
20.            ps_a[i + 1] = ps_a[i] + 1;
21.        else
22.            ps_a[i + 1] = ps_a[i];
23. 
24.        if (S[i] == 'b')
25.            ps_b[i + 1] = ps_b[i] + 1;
26.        else
27.            ps_b[i + 1] = ps_b[i];
28.    }
29. 
30.    ll ans = 0;
31.    rep(l, 0, N)
32.    {
33.        auto itr_a = lower_bound(ps_a.begin() + (l + 1), ps_a.end(), ps_a[l] + A);
34.        auto itr_b = lower_bound(ps_b.begin() + (l + 1), ps_b.end(), ps_b[l] + B);
35. 
36.        int r_a = distance(ps_a.begin(), itr_a) - 1;
37.        int r_b = (distance(ps_b.begin(), itr_b) - 1) - 1;
38. 
39.        if (r_a <= r_b)
40.            ans += r_b - r_a + 1;
41.    }
42. 
43.    cout << ans << endl;
44. 
45.    return 0;
46.}

Submission #70603378 - AtCoder Beginner Contest 430

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

AtCoder

実装時間: 10 分

コメント

尺取り法でも解けるようだが、累積和 + 二分探索の方が個人的には実装しやすかった。