Submission #25825782

---

Source Code Expand

import sys
import numpy as np
from numba import njit, i8

MOD = 998_244_353

@njit
def get_sum(bit, i):
    """sum on [0, i)"""
    s = 0
    while i:
        s += bit[i]
        i -= i & -i
    return s

@njit
def add(bit, i, x=1):
    assert 0 <= i < len(bit) - 1
    i += 1
    while i < len(bit):
        bit[i] += x
        i += i & -i

@njit((i8[:], i8[:, :]), cache=True)
def main(A, query):
    # 座標圧縮
    X = np.unique(np.concatenate((A, query[:, 1])))
    A = np.searchsorted(X, A)
    query[:, 1] = np.searchsorted(X, query[:, 1])
    
    # 個数と和を fenwick tree で持つ
    bit_c = np.zeros(len(X) + 10, np.int64)
    bit_s = np.zeros(len(X) + 10, np.int64)    
    total_c = 0
    total_s = 0
    ANS = 0
    
    def upd(i, k):
        nonlocal ANS, total_c, total_s
        # 多重集合に X[i] を k 個追加する。
        left_c = get_sum(bit_c, i)
        left_s = get_sum(bit_s, i)
        right_c = total_c - left_c
        right_s = total_s - left_s
        
        diff_sum = left_c * X[i] - left_s
        diff_sum += right_s - right_c * X[i]
        # 2 で割る
        if diff_sum & 1:
            diff_sum += MOD
        diff_sum //= 2
        
        ANS = (ANS + k * diff_sum) % MOD
        add(bit_c, i, k)
        add(bit_s, i, X[i] * k)
        total_c += k
        total_s += X[i] * k
        
    for i in A:
        upd(i, 1)

    for q in range(len(query)):
        x, y = query[q]
        x -= 1
        upd(A[x], -1)
        A[x] = y
        upd(A[x], 1)
        print(ANS)

N, Q = map(int, sys.stdin.readline().split())
A = np.array(sys.stdin.readline().split(), np.int64)
query = np.array(sys.stdin.read().split(), np.int64).reshape(Q, 2)

main(A, query)

Submission Info

Judge Result