python~atcoder ABC131 D問題

今回はAtcoderの問題に挑戦します。ABCのコンテストのD問題です。他のD問題に比べて比較的優しく、すっきり解けて気持ちよかったので投稿します。問題文は以下のようになっています。

問題文

AtCoder王国の王立問題工房でABC管理官の座に就いたキザハシ君は、浮かれるあまり仕事を引き受けすぎてしまいました。

現在の時刻は 0です。キザハシ君は 1 から N までの番号が振られた N 件の仕事を持っています。

キザハシ君が仕事 iを終えるには Ai単位時間かかります。また、仕事 i の〆切は時刻 Bi であり、これまでに仕事を終わらせる必要があります。時刻 Bi ちょうどに仕事 i を終わらせてもかまいません。

キザハシ君は 2 件以上の仕事を同時にすることはできませんが、ある仕事を終わらせた直後に別の仕事を始めることはできます。

キザハシ君はすべての仕事を〆切までに終わらせることができるでしょうか。可能ならば Yes、不可能ならば No を出力してください。

制約

• 入力はすべて整数
• 1≤N≤2×10⁵
• 1≤Ai,Bi≤10⁹(1≤i≤N)

入力

入力は以下の形式で標準入力から与えられます。

N
A1 B1
..
..
..
AN BN

出力

全ての仕事を〆切までに終わらせることが可能ならば Yes、不可能ならば No を出力してください。

この問題に対する回答コードは以下になります。

import sys
from operator import itemgetter

input = sys.stdin.readline
N = int(input())
AB = [list(map(int,input().split()))for _ in range(N)]
AB.sort(key=itemgetter(1))
s = 0
ans = "Yes"
for a, b in AB:
  s += a
  if s > b:
    ans = "No"
    break
print(ans)

問題を解くにはまず、入力A,BのBを基準にソートして、締め切りの早い順にします。そしてこの仕事の早い順にAの時間かけて仕事を指摘、B以内に終わっているかを確認していきます。もし途中でBを超えていれば間に合わせられない、最後まで間に合っていれば、間に合わせられることになります。オーダは入力にO(n)、pythonのソートにO(nlogn)、最後のfor文でO(n)なので、全体でO(2n + nlogn)であり、今回はｎは2×10⁵なのトータルで10⁶くらいで、何とか間に合いそうです。随所に高速化を入れれば行けるでしょう。

まず、pythonの入力関数input()ではなく、sys.stdin.readline()を使うことで、高速化しています。

次に[list(map(int,input().split()))for _ in range(N)]のように、繰り返しの入力に対しては内包表記を使うことで、高速化をしています。

AB.sort(key=itemgetter(1))のようにitemgetter()を使って、ABのリストの２列目（インデックスは１）を基準にソートを行っております。

こんな感じで高速化を入れることで、無事通過できました。