Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- class SortedList:
- def __init__(self, iterable=[], _load=200):
- """Initialize sorted list instance."""
- values = sorted(iterable)
- self._len = _len = len(values)
- self._load = _load
- self._lists = _lists = [values[i:i + _load] for i in range(0, _len, _load)]
- self._list_lens = [len(_list) for _list in _lists]
- self._mins = [_list[0] for _list in _lists]
- self._fen_tree = []
- self._rebuild = True
- def _fen_build(self):
- """Build a fenwick tree instance."""
- self._fen_tree[:] = self._list_lens
- _fen_tree = self._fen_tree
- for i in range(len(_fen_tree)):
- if i | i + 1 < len(_fen_tree):
- _fen_tree[i | i + 1] += _fen_tree[i]
- self._rebuild = False
- def _fen_update(self, index, value):
- """Update `fen_tree[index] += value`."""
- if not self._rebuild:
- _fen_tree = self._fen_tree
- while index < len(_fen_tree):
- _fen_tree[index] += value
- index |= index + 1
- def _fen_query(self, end):
- """Return `sum(_fen_tree[:end])`."""
- if self._rebuild:
- self._fen_build()
- _fen_tree = self._fen_tree
- x = 0
- while end:
- x += _fen_tree[end - 1]
- end &= end - 1
- return x
- def _fen_findkth(self, k):
- """Return a pair of (the largest `idx` such that `sum(_fen_tree[:idx]) <= k`, `k - sum(_fen_tree[:idx])`)."""
- _list_lens = self._list_lens
- if k < _list_lens[0]:
- return 0, k
- if k >= self._len - _list_lens[-1]:
- return len(_list_lens) - 1, k + _list_lens[-1] - self._len
- if self._rebuild:
- self._fen_build()
- _fen_tree = self._fen_tree
- idx = -1
- for d in reversed(range(len(_fen_tree).bit_length())):
- right_idx = idx + (1 << d)
- if right_idx < len(_fen_tree) and k >= _fen_tree[right_idx]:
- idx = right_idx
- k -= _fen_tree[idx]
- return idx + 1, k
- def _delete(self, pos, idx):
- """Delete value at the given `(pos, idx)`."""
- _lists = self._lists
- _mins = self._mins
- _list_lens = self._list_lens
- self._len -= 1
- self._fen_update(pos, -1)
- del _lists[pos][idx]
- _list_lens[pos] -= 1
- if _list_lens[pos]:
- _mins[pos] = _lists[pos][0]
- else:
- del _lists[pos]
- del _list_lens[pos]
- del _mins[pos]
- self._rebuild = True
- def _loc_left(self, value):
- """Return an index pair that corresponds to the first position of `value` in the sorted list."""
- if not self._len:
- return 0, 0
- _lists = self._lists
- _mins = self._mins
- lo, pos = -1, len(_lists) - 1
- while lo + 1 < pos:
- mi = (lo + pos) >> 1
- if value <= _mins[mi]:
- pos = mi
- else:
- lo = mi
- if pos and value <= _lists[pos - 1][-1]:
- pos -= 1
- _list = _lists[pos]
- lo, idx = -1, len(_list)
- while lo + 1 < idx:
- mi = (lo + idx) >> 1
- if value <= _list[mi]:
- idx = mi
- else:
- lo = mi
- return pos, idx
- def _loc_right(self, value):
- """Return an index pair that corresponds to the last position of `value` in the sorted list."""
- if not self._len:
- return 0, 0
- _lists = self._lists
- _mins = self._mins
- pos, hi = 0, len(_lists)
- while pos + 1 < hi:
- mi = (pos + hi) >> 1
- if value < _mins[mi]:
- hi = mi
- else:
- pos = mi
- _list = _lists[pos]
- lo, idx = -1, len(_list)
- while lo + 1 < idx:
- mi = (lo + idx) >> 1
- if value < _list[mi]:
- idx = mi
- else:
- lo = mi
- return pos, idx
- def add(self, value):
- """Add `value` to sorted list."""
- _load = self._load
- _lists = self._lists
- _mins = self._mins
- _list_lens = self._list_lens
- self._len += 1
- if _lists:
- pos, idx = self._loc_right(value)
- self._fen_update(pos, 1)
- _list = _lists[pos]
- _list.insert(idx, value)
- _list_lens[pos] += 1
- _mins[pos] = _list[0]
- if _load + _load < len(_list):
- _lists.insert(pos + 1, _list[_load:])
- _list_lens.insert(pos + 1, len(_list) - _load)
- _mins.insert(pos + 1, _list[_load])
- _list_lens[pos] = _load
- del _list[_load:]
- self._rebuild = True
- else:
- _lists.append([value])
- _mins.append(value)
- _list_lens.append(1)
- self._rebuild = True
- def discard(self, value):
- """Remove `value` from sorted list if it is a member."""
- _lists = self._lists
- if _lists:
- pos, idx = self._loc_right(value)
- if idx and _lists[pos][idx - 1] == value:
- self._delete(pos, idx - 1)
- def remove(self, value):
- """Remove `value` from sorted list; `value` must be a member."""
- _len = self._len
- self.discard(value)
- if _len == self._len:
- raise ValueError('{0!r} not in list'.format(value))
- def pop(self, index=-1):
- """Remove and return value at `index` in sorted list."""
- pos, idx = self._fen_findkth(self._len + index if index < 0 else index)
- value = self._lists[pos][idx]
- self._delete(pos, idx)
- return value
- def bisect_left(self, value):
- """Return the first index to insert `value` in the sorted list."""
- pos, idx = self._loc_left(value)
- return self._fen_query(pos) + idx
- def bisect_right(self, value):
- """Return the last index to insert `value` in the sorted list."""
- pos, idx = self._loc_right(value)
- return self._fen_query(pos) + idx
- def count(self, value):
- """Return number of occurrences of `value` in the sorted list."""
- return self.bisect_right(value) - self.bisect_left(value)
- def __len__(self):
- """Return the size of the sorted list."""
- return self._len
- def __getitem__(self, index):
- """Lookup value at `index` in sorted list."""
- pos, idx = self._fen_findkth(self._len + index if index < 0 else index)
- return self._lists[pos][idx]
- def __delitem__(self, index):
- """Remove value at `index` from sorted list."""
- pos, idx = self._fen_findkth(self._len + index if index < 0 else index)
- self._delete(pos, idx)
- def __contains__(self, value):
- """Return true if `value` is an element of the sorted list."""
- _lists = self._lists
- if _lists:
- pos, idx = self._loc_left(value)
- return idx < len(_lists[pos]) and _lists[pos][idx] == value
- return False
- def __iter__(self):
- """Return an iterator over the sorted list."""
- return (value for _list in self._lists for value in _list)
- def __reversed__(self):
- """Return a reverse iterator over the sorted list."""
- return (value for _list in reversed(self._lists) for value in reversed(_list))
- def __repr__(self):
- """Return string representation of sorted list."""
- return 'SortedList({0})'.format(list(self))
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement