ABC271 C - Manga Python解答例

AtCoder Beginner Contest 271 C - MangaをPythonで解きます。

問題

問題文をAtCoderのページより引用します。

問題文

高橋君は全$10^9$巻の漫画『すぬけ君』を読むことにしました。
初め、高橋君は『すぬけ君』の単行本を$N$冊持っています。$i$番目の単行本は$a_i$巻です。

高橋君は『すぬけ君』を読み始める前に限り次の操作を$0$回以上何度でも繰り返せます。

• 現在持っている単行本が$1$冊以下の場合、何もしない。そうでない場合、現在持っている単行本から$2$冊を選んで売り、代わりに好きな巻を選んで$1$冊買う

その後、高橋君は『すぬけ君』を$1$巻、$2$巻、$3$巻、$\ldots$と順番に読みます。ただし、次に読むべき巻を持っていない状態になった場合、(他の巻を持っているかどうかに関わらず)その時点で『すぬけ君』を読むのをやめます。

高橋君が『すぬけ君』を最大で何巻まで読めるかを求めてください。ただし、$1$巻を読めない場合の答えは$0$とします。

制約

• $1 \leq N \leq 3 \times 10^5$
• $1 \leq a_i \leq 10^9$
• 入力はすべて整数

解答例

先頭から順に愚直に求める

読み始めるのは1巻からで、そこから順番に2, 3, ...と読み進めなければなりません。
逆に、後の方の巻をいくら持っていても、先頭の方で抜けがあれば先へは進めません。

そのため、後ろの方の巻を使って穴を埋めながら先頭から見ていく、という方法が最適です。

具体的には、

• 次に見る巻の情報を保持しておく。最初は1から始まる。
• 持っている巻の番号を昇順に並べ、両端キューに全て格納する。
• 両端キューの先頭を見る。
  • 次に見る巻が先頭にあれば、それを1つ取り除く。
  • もし無ければ、後ろの方の2つを取り除く。
• 読める巻が無くなるまでループする。

という感じに解きます。

ただし、先頭の方の巻が重複しているときは注意です。例えば、

5
1 2 2 2 3

という入力が与えられたとき、重複している2巻のうち2つを売って4巻を手に入れれば、合計4冊読むことができますが、上述の処理を行うと3という結果が返ってきます。

これを防ぐには、リストの中で一意な数字を集めてリストの先頭に持ってきて、それを両端キューにすれば良いです。
上の例で言うと、

1 2 2 2 3
↓
1 2 3 2 2

のように並べ替えます。

ソートに$\mathcal{O}(N\log{N})$、前処理に$\mathcal{O}(N)$、本処理に$\mathcal{O}(N)$かかるので、総じて$\mathcal{O}(N\log{N})$で解けます。

実装例

Pythonによる実装例を以下に示します。

c.py

from typing import *
import collections

def main():
    # 入力受け取り
    N: int = int(input())
    # Aは受け取った段階でソートしておく
    A: List[int] = sorted(list(map(int, input().split())))

    # 一意な数値だけを集めるリストB
    B: List[int] = []
    # 重複した余りを集めるリストC
    C: List[int] = []
    # 既に現れた数値を記録しておく集合S
    S: Set[int] = set()
    # 昇順にソートしたリストAを順に見ていく
    for a in A:
        # 既に要素aが現れていたなら、Cに格納する
        if a in S:
            C.append(a)
        else:
            # そうでないなら、Bに格納する
            B.append(a)
            S.add(a)

    # Bの後ろにCを結合する
    B.extend(C)
    # Bを元に両端キューを作る
    D = collections.deque(B)

    # 次に見る巻の数値を保持する変数current
    current: int = 1
    # 答え
    answer: int = 0
    # キューが空になるまで繰り返す
    while D:
        # キューの先頭が、次に見る巻ならばそれを1つ取り除く
        if D[0] == current:
            D.popleft()
            answer += 1
            current += 1
        else:
            # そうでない場合、後ろの2冊を売って目的の巻を入手し、それを読む
            if len(D) >= 2:
                D.pop()
                D.pop()
                answer += 1
                current += 1
            else:
                # 売る本が無ければそこで終了する
                break

    # 解答出力
    print(answer)


if __name__ == "__main__":
    main()

実際の提出はこちら。