ABC136 C Build Stairs Python 3

問題:https://atcoder.jp/contests/abc136/tasks/abc136_c

与えられた数列の各要素に対して、そのままにするか、数を1減らすか選択していき、うまく単調に増加する階段を作ることができるか判定する問題。一手ずつ見落としなくシミュレーションすれば解ける問題。

コード例

左側から順番に確認していきます。

N=int(input())
H=list(map(int,input().split()))
H[0]-=1
for i in range(1,N-1):
	if H[i]>H[i+1]:
		H[i]-=1
		if H[i]>H[i+1]:
			print('No')
			exit()
		if H[i-1]>H[i]:
			print('No')
			exit()
	elif H[i-1]<=H[i]-1:
		H[i]-=1
print('Yes')

まず、単調非減少の階段を作りたいので、初手は必ず−1しておいた方が有利。なので、3行目ではリストの先頭を−1している。

続いて、リストの中身を2番目の要素から、リストの末尾の一つ手前までfor文でチェックしていく。

まず、i番目の要素とi+1番目の要素の大小関係をチェック。もしi+1番目の方が小さいのであれば、リストのi番目すなわちH[i]を−1する。

次に、再度i番目とi+1番目を比べる。これでもまだ、i+1番目の方が小さいのであれば、条件は不成立なため、”No”を返してプログラムは終了。

もし、上記の条件は満たしていたとしても、H[i]を1減らしたせいでi-1番目の要素との関係が崩れていないか確認する必要がある。H[i-1]がH[i]より大きくなってしまうと数列は減少したことになるので、条件は不成立であり、”No”を返して終了である。

最後に、もしH[i]<=H[i+1]を満たしている状態でもなるべくH[i]を小さくしておいた方が有利なので、H[i]から−1引いても大丈夫なら、
つまり、H[i-1]<=H[i]−1 を満たすならH[i]を−1しておく。

これを最後まで繰り返し、途中でExitすることがなければ、数列は問題の条件を満たしたことになるので、”Yes”を出力して終了だ。

サンプルケースとして、問題文とセットになっているもの以外に下記のものもチェックした方が良いかもしれない。自分はこのケースの処理を見落としていてWAになりました。。。

4
1 2 2 1