diff --git a/fenwick/__init__.py b/fenwick/__init__.py index f385d56..bffb0fc 100644 --- a/fenwick/__init__.py +++ b/fenwick/__init__.py @@ -1,3 +1 @@ from .fenwick import FenwickTree, __version__ - -__all__ = ["FenwickTree", "__version__"] diff --git a/fenwick/fenwick.py b/fenwick/fenwick.py index 6d2e62e..a358467 100644 --- a/fenwick/fenwick.py +++ b/fenwick/fenwick.py @@ -38,7 +38,7 @@ def __len__(self): def prefix_sum(self, stop): """Returns sum of first elements (sum up to *stop*, exclusive).""" if stop <= 0 or stop > self._n: - raise IndexError("index out of range") + raise IndexError() _sum = 0 while stop > 0: _sum += self._v[stop - 1] @@ -48,9 +48,9 @@ def prefix_sum(self, stop): def range_sum(self, start, stop): """Returns sum from start (inclusive) to stop (exclusive).""" if start < 0 or start >= self._n: - raise IndexError("index out of range") + raise IndexError() if stop <= start or stop > self._n: - raise IndexError("index out of range") + raise IndexError() result = self.prefix_sum(stop) if start > 0: result -= self.prefix_sum(start) @@ -72,7 +72,7 @@ def frequencies(self): def add(self, idx, k): """Adds k to idx'th element (0-based indexing).""" if idx < 0 or idx >= self._n: - raise IndexError("index out of range") + raise IndexError() idx += 1 while idx <= self._n: self._v[idx - 1] += k @@ -85,9 +85,10 @@ def __setitem__(self, idx, value): def init(self, frequencies): """Initialize in O(n) with specified frequencies.""" - if len(frequencies) != self._n: - raise ValueError("Length of frequencies must match length of FenwickTree.") - self._v = list(frequencies) + n = len(frequencies) + if n != self._n: + raise ValueError(f'length of frequencies ({n}) does not match length of FenwickTree ({self._n})') + self._v[:] = frequencies for idx in range(1, self._n + 1): parent_idx = idx + (idx & -idx) # parent in update tree if parent_idx <= self._n: