#1151885 (PyPy3) No.3507 RangeSum RangeUpdate RangeSqrt

• 問題: No.3507 RangeSum RangeUpdate RangeSqrt
• ユーザー: n_bitand_n_per_3
• 提出日時: 2026-03-07 18:59:40
• 言語: PyPy3 (7.3.17)
• 結果: RE
• 実行時間: -
• コード長: 4,342 bytes
• コンパイル時間: 338 ms
• コンパイル使用メモリ: 85,268 KB
• 実行使用メモリ: 158,848 KB
• 最終ジャッジ日時: 2026-04-17 19:33:51
• 合計ジャッジ時間: 10,992 ms
• ジャッジサーバーID （参考情報）: judge1_0 / judge2_1

ファイルパターン	結果
sample	RE * 1
other	WA * 2 RE * 27

ソースコード

import sys

# 再帰上限を引き上げる
sys.setrecursionlimit(300000)

def solve():
    # 入力読み込み
    input = sys.stdin.read
    data = input().split()
    iterator = iter(data)
    
    N = int(next(iterator))
    
    # 配列Aの読み込み
    A = []
    for _ in range(N):
        A.append(int(next(iterator)))
        
    Q = int(next(iterator))

    # セグメント木のサイズ決定（2のべき乗）
    size = 1
    while size < N:
        size *= 2
        
    # 木のデータ
    # tree_sum: 区間和
    # tree_max: 区間最大値
    # lazy: 遅延代入値 (-1 は遅延なしを示す)
    tree_sum = [0] * (2 * size)
    tree_max = [0] * (2 * size)
    lazy = [-1] * (2 * size)

    # 初期化 (build)
    # 葉の値をセット
    for i in range(N):
        idx = size + i
        val = A[i]
        tree_sum[idx] = val
        tree_max[idx] = val
    
    # 親を計算
    for i in range(size - 1, 0, -1):
        tree_sum[i] = tree_sum[2*i] + tree_sum[2*i+1]
        tree_max[i] = max(tree_max[2*i], tree_max[2*i+1])

    # 遅延評価を子に伝播させる関数 (push)
    def push(x, lx, rx):
        if lazy[x] == -1:
            return
        
        # 子の範囲の長さ
        mid = (lx + rx) // 2
        len_l = mid - lx
        len_r = rx - mid
        
        val = lazy[x]
        
        # 左の子へ伝播
        lazy[2*x] = val
        tree_sum[2*x] = val * len_l
        tree_max[2*x] = val
        
        # 右の子へ伝播
        lazy[2*x+1] = val
        tree_sum[2*x+1] = val * len_r
        tree_max[2*x+1] = val
        
        # 自身の遅延を解消
        lazy[x] = -1

    # 値をマージして親を更新する関数 (pull)
    def pull(x):
        tree_sum[x] = tree_sum[2*x] + tree_sum[2*x+1]
        tree_max[x] = max(tree_max[2*x], tree_max[2*x+1])

    # クエリ2: 区間代入 update_assign(l, r, v)
    def update_assign(l, r, v, x, lx, rx):
        if lx >= r or rx <= l:
            return
        if lx >= l and rx <= r:
            lazy[x] = v
            tree_sum[x] = v * (rx - lx)
            tree_max[x] = v
            return
        
        push(x, lx, rx)
        mid = (lx + rx) // 2
        update_assign(l, r, v, 2*x, lx, mid)
        update_assign(l, r, v, 2*x+1, mid, rx)
        pull(x)

    # クエリ3: 区間Sqrt update_sqrt(l, r)
    def update_sqrt(l, r, x, lx, rx):
        if lx >= r or rx <= l:
            return
        
        # 枝刈り: 最大値が1以下ならSqrtしても変わらない
        if tree_max[x] <= 1:
            return

        # 区間が完全に含まれていて、かつ一様になっている（lazyがある）場合
        # 代入値を直接Sqrtして高速化
        if lx >= l and rx <= r and lazy[x] != -1:
            val = int(lazy[x] ** 0.5)
            lazy[x] = val
            tree_sum[x] = val * (rx - lx)
            tree_max[x] = val
            return

        # 葉ノードの場合（lazyがない場合でも個別に更新）
        if rx - lx == 1:
            val = int(tree_sum[x] ** 0.5)
            tree_sum[x] = val
            tree_max[x] = val
            return

        push(x, lx, rx)
        mid = (lx + rx) // 2
        update_sqrt(l, r, 2*x, lx, mid)
        update_sqrt(l, r, 2*x+1, mid, rx)
        pull(x)

    # クエリ1: 区間和取得 query_sum(l, r)
    def query_sum(l, r, x, lx, rx):
        if lx >= r or rx <= l:
            return 0
        if lx >= l and rx <= r:
            return tree_sum[x]
        
        push(x, lx, rx)
        mid = (lx + rx) // 2
        s1 = query_sum(l, r, 2*x, lx, mid)
        s2 = query_sum(l, r, 2*x+1, mid, rx)
        return s1 + s2

    # クエリ処理
    output = []
    for _ in range(Q):
        type = int(next(iterator))
        l = int(next(iterator))
        r = int(next(iterator)) # 入力は r-1 までだが、実装は半開区間 [l, r) なのでそのまま使える
        
        if type == 1:
            ans = query_sum(l, r, 1, 0, size)
            output.append(str(ans))
        elif type == 2:
            val = int(next(iterator))
            update_assign(l, r, val, 1, 0, size)
        elif type == 3:
            update_sqrt(l, r, 1, 0, size)
            
    sys.stdout.write('\n'.join(output) + '\n')

if __name__ == '__main__':
    solve()