ベスパリブ

プログラミングを主とした日記・備忘録です。ベスパ持ってないです。

ABC 030: C問題の解説 (Python3)

問題:AtCoder Beginner Contest 030: C - 飛行機乗り

公式解説PDF:https://www.slideshare.net/chokudai/abc030

公式解説動画:なし

解説

二分探索問題です。

以下の解説は公式PDFの解説まんまです。

入力例1を考えます。


初期状態は以下のように、高橋くんは時刻0に空港Aにいます。

時刻 0 1 2 3 4 5 6 7 8 9 10 11 12 13
空港Aの出発時刻 1 5 7
空港Bの出発時刻 3 8 12 13
高橋くん A

次に、高橋くんは出発時刻が最も近い時刻1で空港Aを発ち、時刻X(=2)をかけて、時刻3に空港Bに到着します。

時刻 0 1 2 3 4 5 6 7 8 9 10 11 12 13
空港Aの出発時刻 1 5 7
空港Bの出発時刻 3 8 12 13
高橋くん B

次に、高橋くんは出発時刻が最も近い時刻3で空港Bを発ち、時刻Y(=3)をかけて、時刻6に空港Aに到着します。

時刻 0 1 2 3 4 5 6 7 8 9 10 11 12 13
空港Aの出発時刻 1 5 7
空港Bの出発時刻 3 8 12 13
高橋くん A

これで1往復しました。


次に、高橋くんは出発時刻が最も近い時刻7で空港Aを発ち、時刻X(=2)をかけて、時刻9に空港Bに到着します。

時刻 0 1 2 3 4 5 6 7 8 9 10 11 12 13
空港Aの出発時刻 1 5 7
空港Bの出発時刻 3 8 12 13
高橋くん B

次に、高橋くんは出発時刻が最も近い時刻12で空港Bを発ち、時刻X(=3)をかけて、時刻15に空港Aに到着します。

時刻 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
空港Aの出発時刻 1 5 7
空港Bの出発時刻 3 8 12 13
高橋くん A

これで2往復しました。

これ以上空港Aから出発することはできないので、答えは2往復です。


と、考え方は以上になります。

この問題の本質は「現在時刻に最も近い出発時刻をどう見つけるか」です。

入力例1において、空港Bのリストは [3, 8, 12, 13] です。現在時刻が9のとき、9以上で最も近い値はなにか?

これはリストの中身を線形探索するようなやり方をしたら全体の計算量はおそらく O(N^2)になりTLEします。

このような「ある値aに最も近いものを、リストの中から効率よく探す」といえば二分探索です。二分探索はリストの長さをNとすると、 O(logN)で探索することができ、これはリストの中身を線形探索するような O(N)よりもはるかに高速です。

Python3で二分探索はbisectモジュールを使うと簡単です。簡単な実行例は記事の最後に載せています。

コーナーケース

  • 特になし

実装

# -*- coding:utf-8 -*-
import sys
import bisect

def solve():
    N, M = list(map(int, sys.stdin.readline().split()))
    X, Y = list(map(int, sys.stdin.readline().split()))
    As = list(map(int, sys.stdin.readline().split()))
    Bs = list(map(int, sys.stdin.readline().split()))
    As.sort()
    Bs.sort()

    ans = 0  # 最終的に出力する答え
    t = 0    # 現在時刻
    now_pos = "A"  # 現在の空港
    pre_a_i = 0  # 二分探索の前回のインデックス
    pre_b_i = 0  # 二分探索の前回のインデックス
    while True:
        if now_pos == "A":
            i = bisect.bisect_left(As, t, pre_a_i)
            if i >= N: break
            t = As[i] + X  # 空港Aを発ち、空港Bに到着した時刻t
            pre_a_i = i
            now_pos = "B"
        else:
            i = bisect.bisect_left(Bs, t, pre_b_i)
            if i >= M: break
            t = Bs[i] + Y  # 空港Bを発ち、空港Aに到着した時刻t
            pre_b_i = i
            now_pos = "A"
            ans += 1

    print(ans)


if __name__ == "__main__":
    solve()

高速化のためpre_a_i, pre_b_iで探索範囲を狭めています。なくてもACするかも。

計算量は O(N logN + M logM) か...?(わからん)。

参考

bisectモジュールのきほん

import bisect

a = [10, 20, 20, 30, 40]
a.sort()  # aはソート済みであること

bisect_left

公式ドキュメントから引用

  • ソートされた順序を保ったまま x を a に挿入できる点を探し当てます。
  • リストの中から検索する部分集合を指定するには、パラメータの lo と hi を使います。デフォルトでは、リスト全体が使われます。
  • x がすでに a に含まれている場合、挿入点は既存のどのエントリーよりも前(左)になります。
a = [10, 20, 20, 30, 40]

bisect.bisect_left(a, 11)
# 1

bisect.bisect_left(a, 20)
# 1

bisect.bisect_left(a, 20, lo=2)
# 2

bisect.bisect_left(a, 42)
# 5

bisect.bisect_left(a, 20, lo=100)
# 100

bisect_right (bisect)

基本的にbisect_leftと同じ。

  • どのエントリーよりも後ろ(右)にくるような挿入点を返します。
a = [10, 20, 20, 30, 40]

bisect.bisect_right(a, 11)
# 1

bisect.bisect_right(a, 20)
# 3

bisect.bisect_right(a, 20, lo=2)
# 3

bisect.bisect_right(a, 42)
# 5

bisect.bisect_right(a, 20, lo=100)
# 100