fork download
copy 
  1. # ========= SORTED LIST ==========
  2. class SortedList:
  3.     def __init__(self, iterable=[], _load=200):
  4.         """Initialize sorted list instance."""
  5.         values = sorted(iterable)
  6.         self._len = _len = len(values)
  7.         self._load = _load
  8.         self._lists = _lists = [values[i:i + _load] for i in range(0, _len, _load)]
  9.         self._list_lens = [len(_list) for _list in _lists]
  10.         self._mins = [_list[0] for _list in _lists]
  11.         self._fen_tree = []
  12.         self._rebuild = True
  13.  
  14.     def _fen_build(self):
  15.         """Build a fenwick tree instance."""
  16.         self._fen_tree[:] = self._list_lens
  17.         _fen_tree = self._fen_tree
  18.         for i in range(len(_fen_tree)):
  19.             if i | i + 1 < len(_fen_tree):
  20.                 _fen_tree[i | i + 1] += _fen_tree[i]
  21.         self._rebuild = False
  22.  
  23.     def _fen_update(self, index, value):
  24.         """Update `fen_tree[index] += value`."""
  25.         if not self._rebuild:
  26.             _fen_tree = self._fen_tree
  27.             while index < len(_fen_tree):
  28.                 _fen_tree[index] += value
  29.                 index |= index + 1
  30.  
  31.     def _fen_query(self, end):
  32.         """Return `sum(_fen_tree[:end])`."""
  33.         if self._rebuild:
  34.             self._fen_build()
  35.  
  36.         _fen_tree = self._fen_tree
  37.         x = 0
  38.         while end:
  39.             x += _fen_tree[end - 1]
  40.             end &= end - 1
  41.         return x
  42.  
  43.     def _fen_findkth(self, k):
  44.         """Return a pair of (the largest `idx` such that `sum(_fen_tree[:idx]) <= k`, `k - sum(_fen_tree[:idx])`)."""
  45.         _list_lens = self._list_lens
  46.         if k < _list_lens[0]:
  47.             return 0, k
  48.         if k >= self._len - _list_lens[-1]:
  49.             return len(_list_lens) - 1, k + _list_lens[-1] - self._len
  50.         if self._rebuild:
  51.             self._fen_build()
  52.  
  53.         _fen_tree = self._fen_tree
  54.         idx = -1
  55.         for d in reversed(range(len(_fen_tree).bit_length())):
  56.             right_idx = idx + (1 << d)
  57.             if right_idx < len(_fen_tree) and k >= _fen_tree[right_idx]:
  58.                 idx = right_idx
  59.                 k -= _fen_tree[idx]
  60.         return idx + 1, k
  61.  
  62.     def _delete(self, pos, idx):
  63.         """Delete value at the given `(pos, idx)`."""
  64.         _lists = self._lists
  65.         _mins = self._mins
  66.         _list_lens = self._list_lens
  67.  
  68.         self._len -= 1
  69.         self._fen_update(pos, -1)
  70.         del _lists[pos][idx]
  71.         _list_lens[pos] -= 1
  72.  
  73.         if _list_lens[pos]:
  74.             _mins[pos] = _lists[pos][0]
  75.         else:
  76.             del _lists[pos]
  77.             del _list_lens[pos]
  78.             del _mins[pos]
  79.             self._rebuild = True
  80.  
  81.     def _loc_left(self, value):
  82.         """Return an index pair that corresponds to the first position of `value` in the sorted list."""
  83.         if not self._len:
  84.             return 0, 0
  85.  
  86.         _lists = self._lists
  87.         _mins = self._mins
  88.  
  89.         lo, pos = -1, len(_lists) - 1
  90.         while lo + 1 < pos:
  91.             mi = (lo + pos) >> 1
  92.             if value <= _mins[mi]:
  93.                 pos = mi
  94.             else:
  95.                 lo = mi
  96.  
  97.         if pos and value <= _lists[pos - 1][-1]:
  98.             pos -= 1
  99.  
  100.         _list = _lists[pos]
  101.         lo, idx = -1, len(_list)
  102.         while lo + 1 < idx:
  103.             mi = (lo + idx) >> 1
  104.             if value <= _list[mi]:
  105.                 idx = mi
  106.             else:
  107.                 lo = mi
  108.  
  109.         return pos, idx
  110.  
  111.     def _loc_right(self, value):
  112.         """Return an index pair that corresponds to the last position of `value` in the sorted list."""
  113.         if not self._len:
  114.             return 0, 0
  115.  
  116.         _lists = self._lists
  117.         _mins = self._mins
  118.  
  119.         pos, hi = 0, len(_lists)
  120.         while pos + 1 < hi:
  121.             mi = (pos + hi) >> 1
  122.             if value < _mins[mi]:
  123.                 hi = mi
  124.             else:
  125.                 pos = mi
  126.  
  127.         _list = _lists[pos]
  128.         lo, idx = -1, len(_list)
  129.         while lo + 1 < idx:
  130.             mi = (lo + idx) >> 1
  131.             if value < _list[mi]:
  132.                 idx = mi
  133.             else:
  134.                 lo = mi
  135.  
  136.         return pos, idx
  137.  
  138.     def add(self, value):
  139.         """Add `value` to sorted list."""
  140.         _load = self._load
  141.         _lists = self._lists
  142.         _mins = self._mins
  143.         _list_lens = self._list_lens
  144.  
  145.         self._len += 1
  146.         if _lists:
  147.             pos, idx = self._loc_right(value)
  148.             self._fen_update(pos, 1)
  149.             _list = _lists[pos]
  150.             _list.insert(idx, value)
  151.             _list_lens[pos] += 1
  152.             _mins[pos] = _list[0]
  153.             if _load + _load < len(_list):
  154.                 _lists.insert(pos + 1, _list[_load:])
  155.                 _list_lens.insert(pos + 1, len(_list) - _load)
  156.                 _mins.insert(pos + 1, _list[_load])
  157.                 _list_lens[pos] = _load
  158.                 del _list[_load:]
  159.                 self._rebuild = True
  160.         else:
  161.             _lists.append([value])
  162.             _mins.append(value)
  163.             _list_lens.append(1)
  164.             self._rebuild = True
  165.  
  166.     def discard(self, value):
  167.         """Remove `value` from sorted list if it is a member."""
  168.         _lists = self._lists
  169.         if _lists:
  170.             pos, idx = self._loc_right(value)
  171.             if idx and _lists[pos][idx - 1] == value:
  172.                 self._delete(pos, idx - 1)
  173.  
  174.     def remove(self, value):
  175.         """Remove `value` from sorted list; `value` must be a member."""
  176.         _len = self._len
  177.         self.discard(value)
  178.         if _len == self._len:
  179.             raise ValueError('{0!r} not in list'.format(value))
  180.  
  181.     def pop(self, index=-1):
  182.         """Remove and return value at `index` in sorted list."""
  183.         pos, idx = self._fen_findkth(self._len + index if index < 0 else index)
  184.         value = self._lists[pos][idx]
  185.         self._delete(pos, idx)
  186.         return value
  187.  
  188.     def bisect_left(self, value):
  189.         """Return the first index to insert `value` in the sorted list."""
  190.         pos, idx = self._loc_left(value)
  191.         return self._fen_query(pos) + idx
  192.  
  193.     def bisect_right(self, value):
  194.         """Return the last index to insert `value` in the sorted list."""
  195.         pos, idx = self._loc_right(value)
  196.         return self._fen_query(pos) + idx
  197.  
  198.     def count(self, value):
  199.         """Return number of occurrences of `value` in the sorted list."""
  200.         return self.bisect_right(value) - self.bisect_left(value)
  201.  
  202.     def __len__(self):
  203.         """Return the size of the sorted list."""
  204.         return self._len
  205.  
  206.     def __getitem__(self, index):
  207.         """Lookup value at `index` in sorted list."""
  208.         pos, idx = self._fen_findkth(self._len + index if index < 0 else index)
  209.         return self._lists[pos][idx]
  210.  
  211.     def __delitem__(self, index):
  212.         """Remove value at `index` from sorted list."""
  213.         pos, idx = self._fen_findkth(self._len + index if index < 0 else index)
  214.         self._delete(pos, idx)
  215.  
  216.     def __contains__(self, value):
  217.         """Return true if `value` is an element of the sorted list."""
  218.         _lists = self._lists
  219.         if _lists:
  220.             pos, idx = self._loc_left(value)
  221.             return idx < len(_lists[pos]) and _lists[pos][idx] == value
  222.         return False
  223.  
  224.     def __iter__(self):
  225.         """Return an iterator over the sorted list."""
  226.         return (value for _list in self._lists for value in _list)
  227.  
  228.     def __reversed__(self):
  229.         """Return a reverse iterator over the sorted list."""
  230.         return (value for _list in reversed(self._lists) for value in reversed(_list))
  231.  
  232.     def __repr__(self):
  233.         """Return string representation of sorted list."""
  234.         return 'SortedList({0})'.format(list(self))
  235.  
  236.  
  237. # syntax
  238. s = SortedList()
  239. s.add(1)
  240. s.add(5)
  241. s.add(0)
  242.  
  243.  
Success #stdin #stdout 0.02s 9120KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Standard output is empty
https://ideone.com/RyMNKu
language:
Python 3 (python  3.9.5)
created:
2 years ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!