< uruneaさんのAI生成コード置き場

マージソート木

日時はサーバーが記録したものです。更新すると更新日時が変わるので、不都合がある場合は新しく投稿してください。

class MergeSortTree:
    def __init__(self, A):
        self.n = len(A)
        self.size = 1
        while self.size < self.n:
            self.size <<= 1
        
        self.data = [[] for _ in range(2 * self.size)]
        self.csum = [[] for _ in range(2 * self.size)]
        
        for i in range(self.n):
            self.data[self.size + i] = [A[i]]
            self.csum[self.size + i] = [0, A[i]]
        
        for i in range(self.size - 1, 0, -1):
            left = self.data[2 * i]
            right = self.data[2 * i + 1]
            
            merged = sorted(left + right)
            self.data[i] = merged
            
            cs = [0] * (len(merged) + 1)
            for j in range(len(merged)):
                cs[j+1] = cs[j] + merged[j]
            self.csum[i] = cs

    # [l, r), 0-idx
    def query(self, l, r, x):
        l += self.size
        r += self.size
        res = 0
        while l < r:
            if l & 1:
                idx = bisect.bisect_right(self.data[l], x)
                res += self.csum[l][idx]
                l += 1
            if r & 1:
                r -= 1
                idx = bisect.bisect_right(self.data[r], x)
                res += self.csum[r][idx]
            l >>= 1
            r >>= 1
        return res