【AtCoder解説】PythonでABC172のA,B,C問題を制する！

AtCoder Beginner Contest 172のA,B,C問題を、Python3でなるべく丁寧に解説していきます。

ただ解けるだけの方法ではなく、次の3つのポイントを満たす解法を解説することを目指しています。

• シンプル：余計なことを考えずに済む
• 実装が楽：ミスやバグが減ってうれしい
• 時間がかからない：パフォが上がって、後の問題に残せる時間が増える

目次

ABC172 まとめ
A問題『Calc』
B問題『Minor Change』
C問題『Tsundoku』

ABC172 まとめ

全提出人数: 10137人

パフォーマンス

パフォ	AC	点数	時間	順位(Rated内)
200	AB----	300	7分	7435(7276)位
400	AB----	300	5分	6405(6246)位
600	AB----	300	3分	5234(5076)位
800	ABC---	600	119分	4029(3875)位
1000	AB-D--	700	76分	2899(2750)位
1200	ABCD--	1000	91分	1984(1838)位
1600	ABCD--	1000	26分	834(697)位
2000	ABCDE-	1500	62分	309(199)位

色別の問題AC率

※参加前の表示レート（内部レートではない）、灰に初参加者は含めず

問題	AC者数	灰	茶	緑	水	青
A	10050	99.2%	99.6%	99.7%	99.7%	99.7%
B	9757	95.3%	99.4%	99.7%	99.7%	99.7%
C	3234	6.5%	28.9%	71.7%	95.7%	99.2%
D	3119	9.7%	30.0%	58.6%	82.9%	90.8%
E	470	0.2%	0.9%	3.1%	15.9%	42.6%
F	186	0.0%	0.1%	0.3%	3.0%	16.0%

簡易解説

A問題 (10050人AC)『Calc』

素直に書けばOKです。

B問題 (9757人AC)『Minor Change』

S,Tが何文字が違いか数えればいいです。

C問題 (3234人AC)『Tsundoku』

「A,Bの山で、一番上が小さい方を選び続ける貪欲法」ではダメです。Aの山の上から何冊読むか決めると、Bの山のどこまで読めるか決まるので、Aを0冊～N冊まで全探索します。

D問題 (3119人AC)『Sum of Divisors』[この記事では解説しません]

1～1000万まで、全て約数の数を数えると間に合いません。ある数iがどの整数の約数に含まれるか、という発想をすると、約数ごとの合計を等差数列の和の式で求められます。

E問題 (470人AC)『NEQ』[この記事では解説しません]

包除原理を使って、条件を満たさない組み合わせ数を数え上げます。これをすべての組み合わせ数（mPnの2乗）から引くと、条件を満たす組み合わせ数を求められます。

F問題 (186人AC)『Unfair Nim』[この記事では解説しません]

Nimの勝敗はすべての山のXORが0かどうかで決まります。（検索するとすぐ出てきます）Nimについて必要な知識はそれだけで、ゲームの問題ではないらしいです。

A問題『Calc』

問題ページ：A - Calc

問題の意味

そのまま、 $a + a^2 + a^3$ を出力する問題です。

コード

累乗**の優先順位はどの四則演算（+-*/%//など）の演算子よりも高いので、()でくくって(a ** 2)と書かなくても大丈夫です。

a = int(input())
print(a + a ** 2 + a ** 3)

B問題『Minor Change』

問題ページ：B - Minor Change

問題の意味

長さが同じ2つの文字列 $S, T$ が与えられるので、$S$ を何文字変えると $T$ にできるか求める問題です。

考察：

要するに、$S$ と $T$ は何文字違いかわかればいいです。

実装

前から1文字ずつ見て、違う文字なら答えに1足せばいいです。

range(n)でforループを回してもいいですが、せっかくなのでzip(s, t)を使ってみます。zipを使うと、複数のイテラブル（リスト、タプル、文字列など）から、まとめて要素を取り出すことができます。

入力例1の s = 'cupofcoffee'、t = cupofhotteaで説明します。

zipとforを使ったこのコードを実行すると、

for char_s, char_t in zip(s, t):
    print(char_s, char_t)

こうなります。

c c
u u
p p
o o
f f
c h
o o
f t
f t
e e
e a

今回は文字数が同じなので気にする必要はありませんが、zipは短いほうを基準に止まります。

もし長い方を基準にしたければ、itertoolsモジュールのzip_longest関数を使うといいです。足りないときに埋める値を何も指定しなければNone、指定した場合はその値で埋められます。

コード

s = input()
t = input()

ans = 0
for char_s, char_t in zip(s, t):
    if char_s != char_t:
        ans += 1

print(ans)

C問題『Tsundoku』

問題ページ：C - Tsundoku

コード（bisect_right）
コード（めぐる式二分探索）
コード（尺取法）

問題の意味

2つの本の山A, Bがあって、それぞれの一番上に積んである本のうち、どちらかを読むことができます。読んだ本は山から取り除きます。つまり、選んだ山の上から二番目だった本が新しく一番上になります。

はじめにAの上から $i$ 番目である本を読むのにかかる時間は $A_i$ 分、同じくはじめにBの上から $i$ 番目である本を読むのにかかる時間は $B_i$ 分です。

$K$ 分使って、最大で何冊の本を読めるか求めてください。

考察：

「時間がなくなるまで、読むのにかかる時間が小さい本を選び続ける」方法がすぐ思いつきます。しかし、この方法では最大まで本を読めるとは限りません。

例えば次の場合、この選び方だとAの山から $10, 20, 30$ 分と読んで合計 $3$ 冊になりますが、最初にBの山の $50$ 分を選ぶことで、Bの山から11冊読むことができます。

\begin{align}
K &= 60\\
A &= 10, 20, 30\\
B &= 50, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1
\end{align}

貪欲でダメだと気づくには、例えばこのような入力を考えてみます。

\begin{align}
A &= 10, 20, 30, 5, ...\\
B &= 10, 25, 1, 100, ...
\end{align}

1. 「1番上の本が同じ時間のときはどうする？」
2. 「2番目の本を読むのにかかる時間が短いほうを選べばいい？」
3. 「選ばなかったほうの3番目が1分だったらそっちのほうがよくない？」
4. 「いやでも4番目が100分だったら……」
5. 「貪欲じゃダメそう」

という感じでたどり着けると思います。

Aを何冊目まで読むか決めると、Bを何冊読めるか、自動的に決まる

では、どうすればいいのでしょうか。

この問題は、山を選んだ順番は関係なく「最終的にAとB、それぞれ上から何冊目まで読んだか」だけわかれば、『読んだ本の数』と『かかった時間』が決まります。

そこで、「A, Bをそれぞれ上から何冊読むか決めて、合計が $K$ 分以下になるか判定する」ことを「すべてのA,Bの冊数の組み合わせについて全探索する」と解けそうな気がします。

しかし、$N, M$ はそれぞれ最大20万冊なので、最大20万×20万=400億通り試すことになり、2秒では全く間に合いません。

ここで、「Aを上から何冊読むか」を決めると、Bに使える残り時間は $K - (Aの本を読むのに使った時間の合計)$ 分になります。つまり、Aを決めれば、「Bを上から何冊目まで読めるか」が自動的に決まります。

ですので、「Aを上から $an$ 冊読んだとき、Bを最大何冊読めるか」を効率的に求められれば、「Aを $0 ～ N$ 冊読む場合」を全て試すことで、2秒以内に解けます。

累積和

まず、毎回「Aの上から $an$ 番目までの和」と「Bの上から $bn$ 番目までの和」を計算していては間に合いません。

そこで、あらかじめ『累積和』を計算しておきます。累積和とは、「ある数列の、$1$番目から $k$ 番目までの和」の数列です。

\begin{align}
A &= 1, 2, 3, 10, 20\\
S &= 1, 3, 6, 16, 36
\end{align}

累積和を求めるときは、以下の関係を利用すれば、計算量 $O(N)$ （$N$ に比例した計算量）で求めることができます。

S_{i+1} = S_{i} + A_{i+1}

式で書くと難しそうですが、「1～3番目の和は、1～2番目の和に、3番目を足したもの」「1～4番目の和は、さっき求めた1～3番目の和に、4番目を足したもの」というだけです。

具体的には、次のようなコードを書けばいいです。

a_acc = [0] * n
a_acc[0] = a[0]

for i in range(n-1):
    a_acc[i+1] = a_acc[i] + a[i+1]

しかし実は、itertoolsモジュールのaccumulate関数を使うと累積和を求められるので、自力で実装する必要はありません。

from itertools import accumulate
a_acc = list(accumulate(a))
b_acc = list(accumulate(b))

累積和をあらかじめ求めたので、「Aの上から $an$ 冊目まで読むのにかかる時間」（Bも同様）が一瞬でわかるようになりました。

「Bの上から何冊目まで読めるか？」を効率的に求める、2つの方法

これだけではまだ解けません。

「Aの上から何冊読むか決めたとき、Bの上から何冊目まで読めるか」を求めるとき、「Bの累積和 b_acc をb_acc[0] から順番に見ていく」方法がを考えつくと思います。

for an in range(n+1):
    rem = k - a_acc[an]
    bn = 0
    while b_acc[bn] <= rem:
        bn += 1

この方法は計算量が $O(NM)$ （$N\times{M}$ に比例した計算量）になります。$N, M$ はそれぞれ最大20万なので、やはり間に合いません。

これを効率的に求めるアルゴリズムは2つあります。（各アルゴリズムの詳しい解説は省きます）

• 尺取法
• 二分探索

どちらを使うかですが、二分探索のほうが楽なので、二分探索を使ってください。

通常、尺取法のほうが二分探索よりも高速です。この問題の場合、尺取法は $O(N+M)$ 、二分探索は $O(Nlog{M})$ です。

速いほうが良さそうなので、尺取法にしたくなるかもしれません。ですが、二分探索でも十分間に合います。

尺取法はやること自体は単純で、基本的なアルゴリズムとして紹介されることも多いです。しかし、バグらせずに書くのがとても難しく、青コーダーでも「尺取法がバグる」と言っている人をたびたび見かけます。

一方、二分探索は毎回ほとんどテンプレ通りに書けば動きます。さらにいえば、bisectモジュールを使えば自分で実装する必要すらありません。

「尺取法と二分探索、どちらも使えるなら二分探索を使う」ことを意識しておくといいです。よくある $N = 2\times10^{5}$ なら、 $O(NlogN)$ の二分探索が間に合います。

二分探索の実装（bisect_right）

二分探索を自力で実装してもいいですが、bisectモジュールのbisect.bisect_right関数を使っても楽です。bisect_leftもありますが、今回はこちらではありません。

この関数は、「xをソートされた配列aに、ソートを崩さずに挿入するときのインデックス」を二分探索で求める関数です。

bisect_rightはxがaにすでに含まれているとき、既存のそれらよりも右に来るようなインデックスを返します。bisect_leftは反対に、左に来るインデックスを返します。

簡単にいうと、bisect_rightは 「ある数字『以下』がいくつあるか求める」 関数です。一方、bisect_leftは 「ある数字『未満』がいくつあるか求める」 関数です。

例えば、a = [1, 2, 3, 4]のとき、bisect_left(2, a)は1を返し、bisect.bisect_right(2, a)は2を返します。

今回は、残り時間がちょうどでも読める本が増えるので、『以下』を求めるbisect_rightを使います。

bisect_rightを使って、Bの累積和 b_acc に rem = k - a_acc[an] 分以下が bn 個あるとわかったとします。

それら bn 個の中で最大の要素はb_acc[bn-1] = sum(b[0:bn]) （1冊目からbn冊目までの和）なので、最大でbn冊読めます。したがって、an + bn冊読めることになります。

コード

せっかくなので、3通りのコードを載せておきます。

• bisect_right
• めぐる式二分探索
• 尺取法

コード（bisect_right）

「Aを1冊も読まない」場合があるので、forループで回すa_accの先頭に[0]を追加する必要があります。

一方、b_accには[0]を追加しません。例えば残り0分のとき、「0分以下が1個あるので、Bから1冊読める」となってしまいます。

Aの全探索は $O(N)$ 、Bの二分探索は1回につき $O(logM)$ なので、合わせて $O(NlogM)$ になります。

実行時間はPython3では269ms、PyPy3では193msでした。

from itertools import accumulate
import bisect

n, m, k = map(int, input().split())
a = list(map(int, input().split()))
b = list(map(int, input().split()))

a_acc = [0] + list(accumulate(a))  # aを0冊読む=1冊も読まない場合があるので、[0]を先頭に追加します
b_acc = list(accumulate(b))  # bisect_rightでは、x以下の要素がいくつあるか？を求めるので、こちらに0は入れません

ans = 0
for an in range(n + 1):
    rem = k - a_acc[an]
    if rem < 0:
        # aをこれ以上読めません
        break

    bn = bisect.bisect_right(b_acc, rem)
    ans = max(ans, an + bn)

print(ans)

コード（めぐる式二分探索）

バグりにくいことで有名な、[ok, ng)の範囲を狭めていく『めぐる式二分探索』です。

私は遊んだことはありませんが、あるえっちなゲームに出てくる『因幡めぐる』というキャラの、競プロなりきりアカウント発の二分探索です。なお、中の人はAtCoder社長のchokudaiさんでした。

「bisect_leftとbisect_rightのどっちを使えばいいかわからない」となったときは、考えるよりも自分で『めぐる式にぶたん』を実装したほうが楽かもしれません。

これも計算量は $O(NlogM)$ になります。

実行時間はPython3では1105ms、PyPy3では178msでした。Python3だと若干怪しいですが、PyPy3なら余裕です。（PyPy 7.3.0はPython3.6準拠なので、一部Python 3.8の機能を使えないことに注意してください）

from itertools import accumulate

n, m, k = map(int, input().split())
a = list(map(int, input().split()))
b = list(map(int, input().split()))

a_acc = [0] + list(accumulate(a))  # aを0冊読む=1冊も読まない場合があるので
b_acc = [0] + list(accumulate(b))  # 今度は[0]が必要です

ans = 0
for an in range(n + 1):
    rem = k - a_acc[an]
    if rem < 0:
        break

    # 最初は 0 <= x < m + 1 の範囲で、この範囲を狭めていきます
    ok = 0
    ng = m + 1   
    # ng = ok + 1のとき、ok <= x < ng は ok <= x < ok + 1で、x = okに確定します
    while abs(ng - ok) > 1:
        mid = (ok + ng) // 2
        # bをmid冊読めるか？
        if b_acc[mid] <= rem:
        # 読めるならok
            ok = mid
        else:
        # 読めないならng
            ng = mid

    bn = ok
    ans = max(ans, an + bn)

print(ans)

コード（尺取法）

Aから読む本を1冊増やしたとします。すると、Bから読める本の数は変わらないか、減るかのどちらかです。考えてみると当たり前なのですが、「Aから読む本を1冊増やしたら、Bから読める本もなぜか1冊増えた」ということは絶対に起きません。

そこで、「前回よりAを1冊増やす」→「それでもK分以内ならBはそのままで、K分を超えているなら、K分以内になるまでBを1ずつ減らす」とすると、効率的に計算できます。 これが尺取法です。

Aの全探索は $O(N)$ 、Bは最悪でも $M$ から $0$ まで $1$ ずつ減るだけなので $O(M)$ で、合わせると $O(N+M)$ です。（※累積和の計算も $O(N+M)$ です）

この尺取法のコードは単純に見えますが、非常にバグらせやすいので、使わなくて済むなら使わない方がいいです。（尺取法を使わざるを得ない問題もあるので、練習はしておくといいかもしれません）

実行時間はPython3では242ms、PyPy3では154msでした。

from itertools import accumulate

n, m, k = map(int, input().split())
a = list(map(int, input().split()))
b = list(map(int, input().split()))

a_acc = [0] + list(accumulate(a))
b_acc = [0] + list((accumulate(b)))

ans = 0
bn = m

for an in range(n + 1):
    rem = k - a_acc[an]
    if rem < 0:
        break

    # rem >= 0 かつ、b_acc[0] = 0なので、bnはどんなに減っても0で必ず止まります。マイナスになってバグることはありません
    while b_acc[bn] > rem:
        bn -= 1
    ans = max(ans, an + bn)

print(ans)