From fa44c7997ba20b2ec938dd1955b9ad002f3df3a9 Mon Sep 17 00:00:00 2001 From: Mukundan314 Date: Wed, 6 Nov 2024 15:28:57 +0530 Subject: [PATCH] small golf and memory reduction in segment tree --- pyrival/data_structures/SegmentTree.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/pyrival/data_structures/SegmentTree.py b/pyrival/data_structures/SegmentTree.py index 7555ee7..c14bedd 100644 --- a/pyrival/data_structures/SegmentTree.py +++ b/pyrival/data_structures/SegmentTree.py @@ -3,22 +3,21 @@ def __init__(self, data, default=0, func=max): """initialize the segment tree with data""" self._default = default self._func = func - self._len = len(data) - self._size = _size = 1 << (self._len - 1).bit_length() + self._len = _len = len(data) - self.data = [default] * (2 * _size) - self.data[_size:_size + self._len] = data - for i in reversed(range(_size)): + self.data = [default] * (2 * _len) + self.data[_len:_len + _len] = data + for i in reversed(range(1, _len)): self.data[i] = func(self.data[i + i], self.data[i + i + 1]) def __delitem__(self, idx): self[idx] = self._default def __getitem__(self, idx): - return self.data[idx + self._size] + return self.data[idx + self._len] def __setitem__(self, idx, value): - idx += self._size + idx += self._len self.data[idx] = value idx >>= 1 while idx: @@ -30,8 +29,8 @@ def __len__(self): def query(self, start, stop): """func of data[start, stop)""" - start += self._size - stop += self._size + start += self._len + stop += self._len res_left = res_right = self._default while start < stop: