|
@@ -31,6 +31,8 @@ from typing import (
|
|
|
overload,
|
|
|
)
|
|
|
|
|
|
+from sortedcontainers import SortedList
|
|
|
+
|
|
|
from hypothesis.errors import HypothesisWarning
|
|
|
|
|
|
ARRAY_CODES = ["B", "H", "I", "L", "Q", "O"]
|
|
@@ -199,27 +201,36 @@ class LazySequenceCopy:
|
|
|
in O(1) time. The full list API is not supported yet but there's no reason
|
|
|
in principle it couldn't be."""
|
|
|
|
|
|
- __mask: Optional[Dict[int, int]]
|
|
|
-
|
|
|
def __init__(self, values: Sequence[int]):
|
|
|
self.__values = values
|
|
|
self.__len = len(values)
|
|
|
- self.__mask = None
|
|
|
+ self.__mask: Optional[Dict[int, int]] = None
|
|
|
+ self.__popped_indices: Optional[SortedList] = None
|
|
|
|
|
|
def __len__(self) -> int:
|
|
|
- return self.__len
|
|
|
+ if self.__popped_indices is None:
|
|
|
+ return self.__len
|
|
|
+ return self.__len - len(self.__popped_indices)
|
|
|
|
|
|
- def pop(self) -> int:
|
|
|
+ def pop(self, i: int = -1) -> int:
|
|
|
if len(self) == 0:
|
|
|
raise IndexError("Cannot pop from empty list")
|
|
|
- result = self[-1]
|
|
|
- self.__len -= 1
|
|
|
+ i = self.__underlying_index(i)
|
|
|
+
|
|
|
+ v = None
|
|
|
if self.__mask is not None:
|
|
|
- self.__mask.pop(self.__len, None)
|
|
|
- return result
|
|
|
+ v = self.__mask.pop(i, None)
|
|
|
+ if v is None:
|
|
|
+ v = self.__values[i]
|
|
|
+
|
|
|
+ if self.__popped_indices is None:
|
|
|
+ self.__popped_indices = SortedList()
|
|
|
+ self.__popped_indices.add(i)
|
|
|
+ return v
|
|
|
|
|
|
def __getitem__(self, i: int) -> int:
|
|
|
- i = self.__check_index(i)
|
|
|
+ i = self.__underlying_index(i)
|
|
|
+
|
|
|
default = self.__values[i]
|
|
|
if self.__mask is None:
|
|
|
return default
|
|
@@ -227,18 +238,34 @@ class LazySequenceCopy:
|
|
|
return self.__mask.get(i, default)
|
|
|
|
|
|
def __setitem__(self, i: int, v: int) -> None:
|
|
|
- i = self.__check_index(i)
|
|
|
+ i = self.__underlying_index(i)
|
|
|
if self.__mask is None:
|
|
|
self.__mask = {}
|
|
|
self.__mask[i] = v
|
|
|
|
|
|
- def __check_index(self, i: int) -> int:
|
|
|
+ def __underlying_index(self, i: int) -> int:
|
|
|
n = len(self)
|
|
|
if i < -n or i >= n:
|
|
|
raise IndexError(f"Index {i} out of range [0, {n})")
|
|
|
if i < 0:
|
|
|
i += n
|
|
|
assert 0 <= i < n
|
|
|
+
|
|
|
+ if self.__popped_indices is not None:
|
|
|
+ # given an index i in the popped representation of the list, compute
|
|
|
+ # its corresponding index in the underlying list. given
|
|
|
+ # l = [1, 4, 2, 10, 188]
|
|
|
+ # l.pop(3)
|
|
|
+ # l.pop(1)
|
|
|
+ # assert l == [1, 2, 188]
|
|
|
+ #
|
|
|
+ # we want l[i] == self.__values[f(i)], where f is this function.
|
|
|
+ assert len(self.__popped_indices) <= len(self.__values)
|
|
|
+
|
|
|
+ for idx in self.__popped_indices:
|
|
|
+ if idx > i:
|
|
|
+ break
|
|
|
+ i += 1
|
|
|
return i
|
|
|
|
|
|
|
|
@@ -345,14 +372,6 @@ def find_integer(f: Callable[[int], bool]) -> int:
|
|
|
return lo
|
|
|
|
|
|
|
|
|
-def pop_random(random: Random, seq: LazySequenceCopy) -> int:
|
|
|
- """Remove and return a random element of seq. This runs in O(1) but leaves
|
|
|
- the sequence in an arbitrary order."""
|
|
|
- i = random.randrange(0, len(seq))
|
|
|
- swap(seq, i, len(seq) - 1)
|
|
|
- return seq.pop()
|
|
|
-
|
|
|
-
|
|
|
class NotFound(Exception):
|
|
|
pass
|
|
|
|