fork download
  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
stdin
Standard input is empty
stdout
Standard output is empty