fenwick_tree_RAQ

ソースコード

from titan_pylib.data_structures.fenwick_tree.fenwick_tree_RAQ import FenwickTreeRAQ

view on github

展開済みコード

  1# from titan_pylib.data_structures.fenwick_tree.fenwick_tree_RAQ import FenwickTreeRAQ
  2from typing import Iterable, Sequence, Union
  3
  4
  5class FenwickTreeRAQ:
  6    """区間加算/区間和クエリができます。。"""
  7
  8    def __init__(self, n_or_a: Union[Iterable[int], int]):
  9        """構築します。
 10        :math:`O(n)` です。
 11
 12        Args:
 13          n_or_a (Union[Iterable[int], int]): 構築元のものです。
 14        """
 15        if isinstance(n_or_a, int):
 16            self.n = n_or_a
 17            self.bit0 = [0] * (n_or_a + 2)
 18            self.bit1 = [0] * (n_or_a + 2)
 19            self.bit_size = self.n + 1
 20        else:
 21            if not isinstance(n_or_a, Sequence):
 22                n_or_a = list(n_or_a)
 23            self.n = len(n_or_a)
 24            self.bit0 = [0] * (self.n + 2)
 25            self.bit1 = [0] * (self.n + 2)
 26            self.bit_size = self.n + 1
 27            for i, e in enumerate(n_or_a):
 28                self.add_range(i, i + 1, e)
 29
 30    def __add(self, bit: list[int], k: int, x: int) -> None:
 31        k += 1
 32        while k <= self.bit_size:
 33            bit[k] += x
 34            k += k & -k
 35
 36    def __pref(self, bit: list[int], r: int) -> int:
 37        ret = 0
 38        while r > 0:
 39            ret += bit[r]
 40            r -= r & -r
 41        return ret
 42
 43    def add(self, k: int, x: int) -> None:
 44        """``k`` 番目に ``x`` を加算します。
 45        :math:`O(\\log{n})` です。
 46
 47        Args:
 48          k (int):
 49          x (int):
 50        """
 51        assert (
 52            0 <= k < self.n
 53        ), f"IndexError: {self.__class__.__name__}.add({k}, {x}), n={self.n}"
 54        self.add_range(k, k + 1, x)
 55
 56    def add_range(self, l: int, r: int, x: int) -> None:
 57        """区間 ``[l, r)`` に ``x`` を加算します。
 58        :math:`O(\\log{n})` です。
 59
 60        Args:
 61          l (int):
 62          r (int):
 63          x (int):
 64        """
 65        assert (
 66            0 <= l <= r <= self.n
 67        ), f"IndexError: {self.__class__.__name__}.add_range({l}, {r}, {x}), l={l},r={r},n={self.n}"
 68        self.__add(self.bit0, l, -x * l)
 69        self.__add(self.bit0, r, x * r)
 70        self.__add(self.bit1, l, x)
 71        self.__add(self.bit1, r, -x)
 72
 73    def sum(self, l: int, r: int) -> int:
 74        """区間 ``[l, r)`` の総和を返します。
 75        :math:`O(\\log{n})` です。
 76
 77        Args:
 78          l (int):
 79          r (int):
 80
 81        Returns:
 82          int:
 83        """
 84        assert (
 85            0 <= l <= r <= self.n
 86        ), f"IndexError: {self.__class__.__name__}.sum({l}, {r}), l={l},r={r},n={self.n}"
 87        return (
 88            self.__pref(self.bit0, r)
 89            + r * self.__pref(self.bit1, r)
 90            - self.__pref(self.bit0, l)
 91            - l * self.__pref(self.bit1, l)
 92        )
 93
 94    def tolist(self) -> list[int]:
 95        """``list`` にして返します。
 96
 97        Returns:
 98          list[int]:
 99        """
100        return [self.sum(i, i + 1) for i in range(self.n)]
101
102    def __getitem__(self, k: int) -> int:
103        """``k`` 番目の値を返します。
104        ``sum(k, k+1)`` と等価です。
105        :math:`O(\\log{n})` です。
106
107        Args:
108          k (int):
109
110        Returns:
111          int:
112        """
113        assert (
114            0 <= k < self.n
115        ), f"IndexError: {self.__class__.__name__}[{k}], n={self.n}"
116        return self.sum(k, k + 1)
117
118    def __str__(self):
119        return str(self.tolist())
120
121    def __repr__(self):
122        return f"{self.__class__.__name__}({self})"

仕様

class FenwickTreeRAQ(n_or_a: Iterable[int] | int)[source]

Bases: object

区間加算/区間和クエリができます。。

__getitem__(k: int) → int[source]

k 番目の値を返します。 sum(k, k+1) と等価です。 O(log⁡n) です。

Parameters:

k (int)

Return type:

int

add(k: int, x: int) → None[source]

k 番目に x を加算します。 O(log⁡n) です。

Parameters:

  • k (int)

  • x (int)

add_range(l: int, r: int, x: int) → None[source]

区間 [l, r) に x を加算します。 O(log⁡n) です。

Parameters:

  • l (int)

  • r (int)

  • x (int)

sum(l: int, r: int) → int[source]

区間 [l, r) の総和を返します。 O(log⁡n) です。

Parameters:

  • l (int)

  • r (int)

Return type:

int

tolist() → list[int][source]

list にして返します。

Return type:

list[int]