AtCoder-ABC183 D - Water Heater【Python解答例】 - ebisukeプログラミング初心者脱出黙示録

コンテストに出たときは解けなかったAtCoder Beginner Contest183のD問題に挑戦します。

いもす法を使って解いていきたいと思います。
ebisuke33.hatenablog.com

AtCoder Beginner Contest183 D - Water Heater

問題文

給湯器が1つあり、毎分Wリットルのお湯を供給することができます。
N人の人がいます。i番目の人は時刻SiからTiまでの間 (時刻Tiちょうどを除く)、この湯沸かし器で沸かしたお湯を毎分Piリットル使おうと計画しています。お湯はすぐ冷めてしまうので、溜めておくことはできません。
すべての人に計画通りにお湯を供給することはできますか？

制約

・1≤N≤2×10^5
・0≤Si

解答例

n, w = map(int,input().strip().split())

a = []
for i in range(n):
    array = list(map(int, input().strip().split()))
    a.append(array)

t = [0] * 200100 # (※1)すべての要素が0の配列　tの要素数は2×10^5以上

for i in range(n):
    t[a[i][0]] += a[i][2]  # (※2)時刻Si(=a[i][0])にPi(=a[i][2])だけ増加
    t[a[i][1]] -= a[i][2]  # (※3)時刻Ti(=a[i][1])にPi(=a[i][2])だけ減少

can = 1  # (※4)お湯の量が足りるか判定フラグ

for j in range(200050):
    t[j+1] += t[j] # (※5)累積和の計算
    if t[j] > w:
        can = 0 # (※6)時刻tの使用量t[i]がお湯の供給量Wを超えたらcanフラグを0に
        break

if can == 1:
    print("Yes")
else:
    print("No")

解説

いもす法では※1のように要素がすべて0の配列を用意します。
この問題ではTiは2×10^5という制約がありますので、用意した配列tの要素数は2×10^5より大きい200100としました。
※2ではi番目の人がお湯を使い始める時刻Si(=a[i][0])においてお湯の使用量Pi(=a[i][2])分だけ配列tの要素を増加させています。
※3では反対にi番目の人がお湯を使い終わる時刻Ti(=a[i][1])で、お湯の使用量Pi(=a[i][2])分だけ配列tの要素を減少させています。
次に※4では判定する準備としてcanフラグを作成し、1を代入しました。判定ループを抜けても1のままならお湯の供給量は足りることになります。
※2と※3でお湯の使い初めと使い終わりの情報が入った配列が完成しましたので、※5で累積和を計算していきます。
これで配列tの要素は、その時刻におけるN人のお湯の使用量を表します。
お湯の供給量Wを配列tの要素が越えなければcanフラグは1のままで、超える場合はcanフラグに0が代入されループを抜けます。
ループを抜けた後にcanフラグの値に合わせて出力することでACとなりました。

NGパターン

次のコードは実際に私がコンテストに出たときに提出したものでTLEとなったものです。

n, w = map(int,input().strip().split())

a = []
for i in range(n):
    array = list(map(int, input().strip().split()))
    a.append(array)

can = 1

for i in range(201000):
    sum_p = 0
    for j in range(n):
        if a[j][0] <= i and a[j][1] > i:
            sum_p += a[j][2]
    if sum_p > w:
        can = 0

if can == 1:
    print("Yes")
else:
    print("No")