사용자 도구

사이트 도구


ps:teflib:fenwicktree

fenwicktree.py

imports and globals

from typing import Iterable, Union

FenwickTree

코드

# N FenwickTree
# I {"version": "1.1", "typing": ["Iterable", "Union"]}
class FenwickTree:
    """Fenwick tree for sum operation."""

    def __init__(self, nums_or_size: Union[Iterable[float], int]):
        if isinstance(nums_or_size, int):
            self._size = nums_or_size
            self._arr = [0] * nums_or_size
            self._tree = [0] * nums_or_size
        else:
            self._arr = list(nums_or_size)
            self._size = len(self._arr)
            self._tree = self._arr[:]
            for i, num in enumerate(self._tree):
                if i | (i + 1) < self._size:
                    self._tree[i | (i + 1)] += num

    def update(self, pos: int, num: float):
        self._arr[pos] += num
        while pos < self._size:
            self._tree[pos] += num
            pos |= pos + 1

    def set(self, pos: int, num: float):
        self.update(pos, num - self._arr[pos])

    def query(self, beg: int, end: int) -> float:
        res = 0
        i = end - 1
        while i >= 0:
            res += self._tree[i]
            i = (i & (i + 1)) - 1
        i = beg - 1
        while i >= 0:
            res -= self._tree[i]
            i = (i & (i + 1)) - 1
        return res

설명

이 코드를 사용하는 문제

출처문제 번호Page레벨
BOJ16221모독플래티넘 3
BOJ16978수열과 쿼리 22플래티넘 3

OrderStatisticTree

코드

# N OrderStatisticTree
# I {"version": "1.0", "typing": ["Iterable", "Union"]}
class OrderStatisticTree:
    def __init__(self, counts_or_max_num: Union[Iterable[int], int]):
        if isinstance(counts_or_max_num, int):
            self._size = 1 << ((counts_or_max_num + 1).bit_length())
            self._arr = [0] * self._size
            self._tree = [0] * self._size
        else:
            self._arr = list(counts_or_max_num)
            self._size = 1 << ((len(self._arr) + 1).bit_length())
            self._arr += [0] * (self._size - len(self._arr)) 
            self._tree = self._arr[:]
            for i, num in enumerate(self._tree):
                if i | (i + 1) < self._size:
                    self._tree[i | (i + 1)] += num

    def add(self, num: int, count: int = 1):
        self._arr[num] += count
        while num < self._size:
            self._tree[num] += count
            num |= num + 1

    def size(self) -> int:
        return self._tree[-1]

    def count(self, num: int) -> int:
        return self._arr[num]

    def count_less_than(self, num: int) -> int:
        res = 0
        r = num - 1
        while r >= 0:
            res += self._tree[r]
            r = (r & (r + 1)) - 1
        return res

    def kth(self, k: int) -> int:
        pos = -1
        for i in range(self._size.bit_length() - 1, -1, -1):
            p = 1 << i
            v = self._tree[pos + p]
            if v < k:
                k -= v
                pos += p
        return pos + 1

설명

이 코드를 사용하는 문제

출처문제 번호Page레벨
BOJ11012Egg플래티넘 1
BOJ13537수열과 쿼리 1플래티넘 4
BOJ1517버블 소트골드 1
BOJ15899트리와 색깔플래티넘 2
BOJ2517달리기플래티넘 4
BOJ5012불만 정렬플래티넘 3
BOJ7568덩치실버 5

토론

댓글을 입력하세요:
D X V V Y
 
ps/teflib/fenwicktree.txt · 마지막으로 수정됨: 2021/04/30 09:38 저자 teferi