クルトンのプログラミング教室

Pythonの使い方やPythonを使った競技プログラミングの解法などを解説しています。

Pythonで解くAtCoder(ABC176:C)

こんにちは、クルトンです!
この記事では

AtCoderの「AtCoder Beginner Contest 176」のC問題をPythonを使って解説します。

問題のリンクを載せておきます。

atcoder.jp


解法

「踏み台を込めて身長を比較したとき、自分より前に、自分より背の高い人が存在しない」ことが条件なので、前から順番に調べていくとよいでしょう。

このとき、i+1番目の人よりも前にいる人の中で一番背の高い人を求めるために、i人を毎回調べていると計算量がO(N^2)になるため間に合いません。

このような場合は最も背の高い人の身長を変数に代入して保存しておくと、計算量がO(N)になり間に合います。


では実装していきましょう。


実装

まず入力を受け取ります。

N = int(input())
A = list(map(int,input().split()))


次に、踏み台を込めて一番背の高い人の身長を保存するための変数MAXと踏み台の高さの合計値を表す変数ansをおきます。
このとき、一番前の人の前には人がいないので、MAX変数の初期値は0にしておきます。

MAX = 0
ans = 0


ここからは前から順番に踏み台の高さを計算していきます

for i in range(N):
    ans += max(0, MAX - A[i])
    MAX = max(MAX, A[i])


ここで少しコードの説明をしようと思います。


まず

    ans += max(0, MAX - A[i])

の部分ですが、この部分では

「A[i]の方がMAXよりも小さい場合は踏み台が必要であるため踏み台の高さ(MAX - A[i])を足す」

という作業をmax関数を使って行っています。もちろんif文を使って書いてもいいですが、この方が記述量は少なくなるのでオススメです。


次に

    MAX = max(MAX, A[i])

の部分ですが、この部分では変数MAXを更新しています。

「踏み台の高さを考えなくていいの?」と疑問に思われるかもしれませんが、最も身長の高い人は踏み台に乗っていないので問題ありません。


最後に答えを出力します。

print(ans)

サンプルコード

Pythonで書いた解答を載せておきます。


atcoder.jp


まとめ

今回は

  • 計算量に注意する

という問題でした。

ではまたお会いしましょう!さようなら~