セグ木の勉強にABC343Fを解く

筆者はレート600前後の茶色コーダー

この記事で遅延セグ木に一回挫折したのでまずはセグメント木を理解したい

初めて参加したコンテストABC343のF問題にセグ木で解ける問題があるみたいなのでそれを解いてみる。

実装コード

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

• shakayami/ACL-for-python

import sys
from collections import defaultdict
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 segtree():
    #---中略---
    
def segfunc(a,b):
    va0, ca0, va1, ca1 = a
    vb0, cb0, vb1, cb1 = b
    vc0, cc0, vc1, cc1 = (-1, 0, -2, 0)
    
    if va0 > vb0:
        vc0 = va0
        cc0 = ca0
        vc1 = max(va1, vb0)
        if vc1 == va1:
            cc1 += ca1
        if vc1 == vb0:
            cc1 += cb0
    elif va0 == vb0:
        vc0 = va0
        cc0 = ca0 + cb0
        vc1 = max(va1,vb1)
        if vc1 == va1:
            cc1 += ca1
        if vc1 == vb1:
            cc1 += cb1
    else:
        vc0 = vb0
        cc0 = cb0
        vc1 = max(va0,vb1)
        if vc1 == va0:
            cc1 += ca0
        if vc1 == vb1:
            cc1 += cb1
    c = vc0, cc0, vc1, cc1
    return c

_id = (-1,0,-2,0)
def main():
    N, Q = rLI()
    A = rLI()
    s = segtree([(a,1,-1,0) for a in A],segfunc,_id)
    for _ in range(Q):
        q, *args = rLI()
        if q == 1:
            p,x = args
            s.set(p-1,((x,1,-1,0)))
        elif q == 2:
            l,r = args
            ans = s.prod(l-1,r)
            print(ans[3])
if __name__ == '__main__':
    main()

このコードでAC通った
https://atcoder.jp/contests/abc343/submissions/54885658

最初はsegfuncをこのように実装したがTLEが出た

def segfunc1(a,b):
    c = defaultdict(lambda: 0)
    x = a + b
    m0 = -1
    m1 = -2
    for i in range(4):
        k,v = x[2*i:2*(i+1)]
        c[k] += v
        if k > m0:
            m0, m1 = k, m0
        elif m0 > k > m1:
            m1 = k
    return (m0,c[m0],m1,c[m1])

で解説のC++のコード観ながらsegfuncを組み替えたら上記の実装コードでうごいた。

理由がわからないのでGPTに聞いてみた

質問

segfunc1でTLEがでてsegfuncでACできるのはなんで？

(コードはりつけ)

回答

segfunc1がTLE（Time Limit Exceeded）になる理由は、その関数が時間計算量が大きいためです。
具体的には、segfunc1では、aとbを連結した後に、その結果をループで処理しています。
これは、aとbの長さが大きい場合、計算量が増大し、時間制限を超える可能性があります。

一方、segfuncでは、aとbを個別に処理し、その結果を組み合わせています。
これにより、計算量が大幅に削減され、より効率的な処理が可能となります。
したがって、segfuncはsegfunc1よりも高速に動作し、
同じ入力であっても時間制限を超えることなく結果を得ることができます。
これがsegfunc1がTLEになり、segfuncがAC（Accepted）になる理由です。

小さな集合でもループや連結等の時間かかる処理を内部で大量に行われるから遅くなる感じ？

まとめ

セグ木の勉強にABC343Fを解いた
segfuncの実装で雑にループやタプルの連結をしたのでTLEが出たが
解説にあったコードの条件分岐を並べた実装に変更してみたら無事にACを出せた。