权值线段树

与普通线段树并无其他区别,只不过存储的信息是每个值出现的次数罢了

理解图
image

import sys

input = lambda: sys.stdin.readline()


class Tree:
    def __init__(self, N):
        self.cnt = [0 for _ in range(N)]

    def update(self, root, l, r, x, cnt):
        if l == r:
            self.cnt[root] += cnt
            return
        mid = l + r >> 1
        if x <= mid:
            self.update(root << 1, l, mid, x, cnt)
        else:
            self.update(root << 1 | 1, mid + 1, r, x, cnt)
        self.cnt[root] = self.cnt[root << 1] + self.cnt[root << 1 | 1]

    # 值域在[ql,qr]的数出现的次数
    def query(self, root, l, r, ql, qr):
        if ql > r or qr < l:
            return 0
        if ql <= l and qr >= r:
            return self.cnt[root]
        mid = l + r >> 1
        return self.query(root << 1, l, mid, ql, qr) + self.query(root << 1 | 1, mid + 1, r, ql, qr)

    def find_kth(self, root, l, r, k):
        if l == r:
            return l
        mid = l + r >> 1
        if k <= self.cnt[root << 1]:
            return self.find_kth(root << 1, l, mid,k)
        else:
            return self.find_kth(root << 1 | 1, mid + 1, r,k - self.cnt[root << 1])


# 当值域过大时,可以采用离散化的方法!
N = 10
mytree = Tree(N * 4)
lst = [1, 2, 3, 4, 5, 5, 5, 5, 5]
for x in lst:
    mytree.update(1, 1, N, x, 1)

for i in range(1, N + 1):
    print(mytree.query(1, 1, N, i, i), end=' ')
print()
# 1 1 1 1 5 0 0 0 0 0
print(mytree.find_kth(1, 1, N, 4))
# 4
posted @ 2024-03-26 12:29  gebeng  阅读(24)  评论(0)    收藏  举报