CODE FESTIVAL 2014 予選B D - 登山家

問題原文

atcoder.jp

解法

ぱっと思い浮かぶのが、平衡二分探索木 + が大きいものから見て行って、すでに見たやつに対して二分探索。
とはいえ当方はPythonなので＼(^o^)／になる。

ここから解説に載ってた別の解き方。
まず、「ある山小屋のより東であって高い位置にあるような最も西の山小屋」が求められればその逆も求められるので、片方向についてのみ求めることを考える。

キーとなる性質は、「山を東から見て行って、今見ている山小屋から見える山は一度見たら二度考える必要はない」ということ。
例えば、 5 3 3 3 とあって、 5 からは 3 3 3 が見えるが、一度見たらこれは二度考える必要がない。もう一つ西を見て、仮に 4 5 3 3 3 だった場合には 4 は 5 で止まるし、 6 5 3 3 3 だった場合、 6 からは 3 3 3 が見えることは確定している。

これを踏まえると以下のようなアルゴリズムで、「ある山小屋のより東であって高い位置にあるような最も西の山小屋」が求められる。

1. スタックにこれまで見た山小屋を積んでいくことを考える
2. 今見ている山小屋からスタック先頭の山小屋が見える限りpop
3. 2を満たさなくなった時点でスタックの先頭にあるのが求めたい山小屋
4. いま見た山小屋をスタックの先頭にpush

実装

なんとなく共通化したかった。
番兵をおくと楽です！

N = int(input())
H = [int(input()) for _ in range(N)]


def calc(H):
    ret = [0] * N
    stack = [N]
    for i in reversed(range(N)):
        while H[stack[-1]] <= H[i]:
            stack.pop()
        ret[i] = stack[-1] - i - 1
        stack.append(i)
    return ret


ans_l = calc(H + [float('inf')])
ans_r = calc(H[::-1] + [float('inf')])[::-1]
for al, ar in zip(ans_l, ans_r):
    print(al + ar)

感想

先にぱっと解法が見えて、それがどうやら実装困難なことがわかると思考が止まっちゃうな。。。
勉強になりました。

atcoder.jp

と同じ雰囲気を感じますね。 RMQ + 二分探索で殴れはしそう。