Created
May 18, 2020 21:14
Revisions
-
vxgmichel created this gist
May 18, 2020 .There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode charactersOriginal file line number Diff line number Diff line change @@ -0,0 +1,195 @@ """ Fenwick tree - performing both fast sum and fast update on a list of number This implementation is mostly compatible with regular list operations. Time complexity table: | Operation | List | Fenwick tree | |-------------|------|--------------| | Instantiate | O(n) | O(n) | | Append | O(1) | O(1) | | Get item | O(1) | O(log k) | | Set item | O(1) | O(log n) | | Sum a slice | O(k) | O(log k) | | Delete item | O(n) | O(n log n) | | Insert | O(n) | O(n log n) | """ from collections.abc import MutableSequence # Base operations def fenwick_sum(data, index): assert 0 <= index < len(data) index += 1 result = 0 while index != 0: result += data[index - 1] index -= index & -index return result def fenwick_add(data, index, value): assert 0 <= index < len(data) index += 1 while index <= len(data): data[index - 1] += value index += index & -index def fenwick_append(data, value): index = len(data) mask = 1 while index & mask: value += data[index - 1] index -= mask mask <<= 1 data.append(value) # Slice operations def fenwick_slice(data, start, stop): stop = min(stop, len(data)) if start >= stop: return [] assert 0 <= start < stop <= len(data) result = [] previous_sum = fenwick_sum(data, start - 1) if start > 0 else 0 for index in range(start, stop): next_sum = fenwick_sum(data, index) result.append(next_sum - previous_sum) previous_sum = next_sum return result def fenwick_sumslice(data, start, stop): stop = min(stop, len(data)) if start >= stop: return 0 assert 0 <= start < stop <= len(data) result = fenwick_sum(data, stop - 1) if start > 0: result -= fenwick_sum(data, start - 1) return result def fenwick_delslice(data, start, stop): stop = min(stop, len(data)) if start >= stop: return assert 0 <= start < stop <= len(data) extra = fenwick_slice(data, stop, len(data)) del data[start:] for value in extra: fenwick_append(data, value) # Access operations def fenwick_get(data, index): if not 0 <= index < len(data): raise IndexError(index) return fenwick_sumslice(data, index, index + 1) def fenwick_set(data, index, value): if not 0 <= index < len(data): raise IndexError(index) current = fenwick_get(data, index) fenwick_add(data, index, value - current) def fenwick_del(data, index): if not 0 <= index < len(data): raise IndexError(index) return fenwick_delslice(data, index, index + 1) def fenwick_insert(data, index, value): assert 0 <= index extra = fenwick_slice(data, index, len(data)) del data[index:] fenwick_append(data, value) for value in extra: fenwick_append(data, value) class Fenwick(MutableSequence): def __init__(self, iterable=()): self._data = [] for value in iterable: fenwick_append(self._data, value) def __len__(self): return len(self._data) def __getitem__(self, index): if not isinstance(index, slice): return fenwick_get(self._data, index) assert index.step in (1, None) return fenwick_slice(self._data, index.start, index.stop) def __setitem__(self, index, value): if not isinstance(index, slice): return fenwick_set(self._data, index, value) indexes = range(*index.indices(len(self._data))) assert len(value) == len(indexes) for index, value in zip(indexes, value): fenwick_set(self._data, index, value) def __delitem__(self, index): if not isinstance(index, slice): return fenwick_del(self._data, index) assert index.step in (1, None) return fenwick_delslice(self._data, index.start, index.stop) def insert(self, index, value): return fenwick_insert(self._data, index, value) def sum(self, start, stop=None): if stop is None: start, stop = 0, start return fenwick_sumslice(self._data, start, stop) def test(): lst = list(range(1, 10)) fen = Fenwick(lst) assert list(fen) == lst lst[3] = 10 fen[3] = 10 assert list(fen) == lst del lst[3] del fen[3] assert list(fen) == lst lst.insert(3, 11) fen.insert(3, 11) assert list(fen) == lst lst.insert(100, 12) fen.insert(100, 12) assert list(fen) == lst lst[1:3] = [13, 14] fen[1:3] = [13, 14] assert list(fen) == lst lst += [15, 16] fen += [15, 16] assert list(fen) == lst del lst[1:3] del fen[1:3] assert list(fen) == lst for i in range(len(lst)): for j in range(len(lst)): assert fen[i:j] == lst[i:j] assert fen.sum(i, j) == sum(lst[i:j])