遅延セグ木でABC359Eを解こうとして諦めた話

筆者は執筆時点レート600弱の茶色コーダ

最近の目標はDPとセグ木をマスターすることで、
先日行われたABC359のE問題は遅延セグ木でも解けると解説にあったので頑張ってコーディングしていく

実装コード

遅延セグ木の実装は以下のものを使用する

• shakayami/ACL-for-python

import sys

def rI(): return int(sys.stdin.readline().rstrip())
def rLI(): return list(map(int,sys.stdin.readline().rstrip().split()))
def rS(): return sys.stdin.readline().rstrip()
def rLS(): return list(sys.stdin.readline().rstrip().split())
def err(*args): print(*args, file=sys.stderr)

class lazy_segtree():
    # ---省略---

INF = 1 << 63
ID = INF

def op(ele1, ele2):
    return ele1 + ele2

def mapping(func, ele):
    if func == ID:
        return ele
    else:
        return func

def composition(func_upper, func_lower):
    if func_upper == ID:
        return func_lower
    else:
        return func_upper

e = 0
id_ = ID
def main():
    N = rI()
    H = [10**9+1]+rLI()
    zeros = [0]*(N+1)
    A=lazy_segtree(zeros[:],op,0,mapping,composition,id_)
    
    ans = zeros[:]
   
    A.set(0,H[1])

    for i in range(1,N+1):
        k = i
        while True:
            j = A.min_left(k,lambda a: a<H[i])
            k = A.min_left(j,lambda a: a<H[i])
            if k==j:break
        A.apply(k,i,H[i])
        ans[i] = A.prod(0,i)+1
        
    print(*ans[1:])
if __name__ == '__main__':
    main()

実行してみたが…

実行してみた結果、サンプルケースが通らない
online judgeを通すと以下のようにWA判定が出る

[INFO] time: 0.035567 sec
[FAILURE] WA
input:
15
748_169_586_329_972_529_432_519_408_587_138_249_656_114_632

output:
749_918_1921_2250_1945_2474_2906_2993_3401_3706_3844_3955_4569_4683_5833

expected:
749_918_1921_2250_4861_5390_5822_6428_6836_7796_7934_8294_10109_10223_11373

遅延セグ木の中身がよくわからないので
デバッグコードを入れてみる

    for i in range(1,N+1):
        k = i
        while True:
            j = A.min_left(k,lambda a: a<H[i])
            k = A.min_left(j,lambda a: a<H[i])
            if k==j:break
        A.apply(k,i,H[i])
+       P_sum[i] = sum([A.get(p) for p in range(i)])+1
        ans[i] = A.prod(0,i)+1
    
+   err(*P_sum[1:])
    print(*ans[1:])
if __name__ == '__main__':
    main()

そしたらこっちは正しい値を出している

[INFO] sample-3
749 918 1921 2250 4861 5390 5822 6428 6836 7796 7934 8294 10109 10223 11373

なのでこっちを出力してみたけど当然TLE
https://atcoder.jp/contests/abc359/submissions/54879855

ここで挫折

諦めた

結局一番簡単なスタックを使った実装でAC通した
https://atcoder.jp/contests/abc359/submissions/54882294

import sys
from collections import deque
def rI(): return int(sys.stdin.readline().rstrip())
def rLI(): return list(map(int,sys.stdin.readline().rstrip().split()))
def rS(): return sys.stdin.readline().rstrip()
def rLS(): return list(sys.stdin.readline().rstrip().split())
def err(*args): print(*args, file=sys.stderr)

def main():
    N = rI()
    H = rLI()
    
    rects = deque()
    rects.append((0,0))
    
    ans = [0] * (N)
    w = 0
    for i,h in enumerate(H):
        v, c = rects[-1]
        if len(rects) == 0:
            rects.append((h,i+1))
        else:
            c_new = 1
            while len(rects) > 0:
                v, c = rects[-1]
                if h < v:
                    break
                rects.pop()
                w -= v*c
                c_new += c
            rects.append((h,c_new))
            w += h*c_new
        ans[i] = w+1
    print(*ans)
    
if __name__ == '__main__':
    main()

まとめ

セグ木も使いこなせていないのにそれより難しい遅延セグ木に手を出すのが間違っていた

まずはセグ木の問題から解けるようにしたい。