class TwoSum:
"""Stores a set of integers and returns the pairs of values that sum to a
specified total.
"""
def __init__(self, startList=None):
"""Initialize set of integers.
"""
self._intSet = SortedSet()
if startList is not None: self.addIntegerList(startList)
def addIntegerList(self, intList):
"""Add integers from list to set
"""
for i in intList:
self.addInteger(i)
def addInteger(self, integer):
"""Add a single integer to set
"""
self._intSet.add(integer)
def findSums(self, totalLow, totalHigh, requireUnique=False):
"""Return set of totals in range for which at least one pair of
integers in set existins such that x + y = total.
Optional parameter for uniqueness will exclude x + x = integer.
"""
# Flip total high and low if passed out of order
if totalLow > totalHigh:
swap = totalLow
totalLow = totalHigh
totalHigh = swap
# Initalize empty found totals set
foundTotals = set()
# Iterate over each integer in the set
for x in self._intSet:
# Determine search range for y
# NOTE: start index is inclusive where stop index is exclusive
# to match behavior of SortedSet.islice()
iStart = self._intSet.bisect_left(totalLow - x)
iStop = self._intSet.bisect_right(totalHigh - x)
# Continue if search range invalid
if iStart >= iStop: continue
# Add each sum x + y to found totals
for y in self._intSet.islice(iStart, iStop):
foundTotals.add(x + y)
# Return found totals set
return foundTotals
def maxSumSubmatrix(self, matrix: List[List[int]], k: int) -> int:
m, n = len(matrix), len(matrix[0])
res = float("-inf")
for row_start in range(m):
prefix_sum = [0] * n
for row_end in range(row_start, m):
row_prefix_sum = 0
for col in range(n):
row_prefix_sum += matrix[row_end][col]
prefix_sum[col] += row_prefix_sum
seen = SortedSet([0])
for num in prefix_sum:
index = seen.bisect_right(num - k - 1)
if index < len(seen):
res = max(res, num - seen[index])
seen.add(num)
return res
def concert_tickets(n, m, tickets, c_price):
tickets.sort()
cnt = Counter(tickets)
t_values = SortedSet(cnt.keys())
t_values.add(-1)
res = ''
for c in c_price:
i = t_values.bisect_right(c) - 1
if t_values[i] <= c:
p = t_values[i]
else:
i = i - 1
p = t_values[i]
if cnt[p] == 0:
res += '-1\n'
else:
res += str(p) + '\n'
cnt[p] -= 1
if cnt[p] == 0:
cnt.pop(p)
t_values.pop(i)
return res
def test_bisect():
temp = SortedSet(range(100), load=7)
assert all(temp.bisect_left(val) == val for val in range(100))
assert all(temp.bisect(val) == val for val in range(100))
assert all(temp.bisect_right(val) == (val + 1) for val in range(100))
import sys
INT_MAX = sys.maxsize
from sortedcontainers import SortedSet
s = SortedSet() # treeset
cmd = list(map(int, input().split()))
result = INT_MAX
for i in range(cmd[0]):
num = int(input())
bigNum = num + cmd[1]
smallNum = num - cmd[1]
bigIdx = s.bisect_left(bigNum)
smallIdx = s.bisect_right(smallNum) - 1
if smallIdx < 0 and bigIdx == len(s):
s.add(num)
continue
elif bigIdx == len(s):
temp = num - s[smallIdx]
elif smallIdx < 0:
temp = s[bigIdx] - num
else:
smallTemp = num - s[smallIdx]
bigTemp = s[bigIdx] - num
temp = min(smallTemp, bigTemp)
if result > temp:
result = temp
s.add(num)
if result == INT_MAX:
def test_bisect():
temp = SortedSet(range(100))
temp._reset(7)
assert all(temp.bisect_left(val) == val for val in range(100))
assert all(temp.bisect(val) == (val + 1) for val in range(100))
assert all(temp.bisect_right(val) == (val + 1) for val in range(100))
class Selection(IMutableGSlice):
def __init__(
self,
universe: slice,
revealed: list = None,
intervals: Iterator = None,
_length: Optional[int] = None # For performance
):
#assert isinstance(universe, slice) # Should universe even be visible/exist?
#assert universe.start == 0
#assert isinstance(universe.stop, int)
#assert universe.stop >= 1 # TODO Do we need this?
self.universe = universe
if intervals is None and revealed is None:
self._intervals = self.revealed2sortedset([slice(0, universe.stop)])
elif intervals is not None:
self._intervals = SortedSet(intervals)
else:
self._intervals = self.revealed2sortedset(revealed)
self._revealed_count = _length if isinstance(_length, int) else Selection._compute_len(self._intervals)
@staticmethod
def revealed2sortedset(revealed: List[Union[tuple, slice]]) -> SortedSet:
""" Converts a list of included pairs to a sorted set of integers in O(n), n = size of @slices.
Every number from every slice is added to the sorted set, except 0.
"""
# 10, [] -> 10, []
# 10, [(0, 10)] -> 10, [10]
# 10, [(0, 7)] -> 10, [7]
# 10, [(7, 10)] -> 10, [7, 10]
# 10, [(3, 7)] -> 10, [3, 7]
# 10, [(0, 3), (7, 10)] -> 10, [3, 7, 10]
# 10, [(0, 1), (2, 3), (4, 5), (6, 7), (8, 9)] -> 10, [1, 2, 3, 4, 5, 6, 7, 8, 9]
try:
#intervals = SortedSet(a for a, _ in revealed).union(b for _, b in revealed)
intervals = SortedSet()
for a, b in revealed:
intervals.add(a)
intervals.add(b)
except TypeError: # slice
intervals = SortedSet(sl.start for sl in revealed).union(sl.stop for sl in revealed)
if 0 in intervals:
intervals.remove(0)
return intervals
@staticmethod
def sortedset2slices(sortedset: SortedSet) -> List[slice]:
""" Converts a sorted set of integers to a list of included slices in O(n), n = size of @sortedset.
If there is an even number of elements in @sortedset, the first slice is formed by the first and second
numbers, the second slice is formed by the third and fourth numbers, and so on.
If there is an odd number of elements in @sortedset, the pair consisting of the number 0 and the first element
in @sortedset becomes the first slice in the output list. The remaining slices, if any, are formed by the
second and third numbers, the fourth and fifth numbers, and so on.
"""
slices = []
if len(sortedset) % 2 == 0:
for i in range(0, len(sortedset), 2):
slices.append(slice(sortedset[i], sortedset[i + 1]))
else:
slices.append(slice(0, sortedset[0]))
for i in range(1, len(sortedset), 2):
slices.append(slice(sortedset[i], sortedset[i + 1]))
return slices
def slices(self) -> List[slice]:
return self.sortedset2slices(self._intervals)
def pairs(self) -> Iterator[Tuple[int, int]]:
if len(self._intervals) % 2 == 0:
return zip(self._intervals[::2], self._intervals[1::2])
return itertools.chain([(0, self._intervals[0])], zip(self._intervals[1::2], self._intervals[2::2]))
def gap_pairs(self) -> Iterator[Tuple[int, int]]:
return self.complement().pairs()
def intervals(self):
return self._intervals
def exclude(self, from_index: Optional[int], to_index: Optional[int]):
original_length = self._revealed_count
if isinstance(from_index, int) and -self.universe.stop <= from_index < 0:
from_index = from_index % self.universe.stop
if isinstance(to_index, int):
if to_index > self.universe.stop:
return self.exclude(from_index, None)
if -self.universe.stop <= to_index < 0:
to_index = to_index % self.universe.stop
assert from_index is None or self.universe.start <= from_index <= self.universe.stop
assert to_index is None or self.universe.start <= to_index <= self.universe.stop
if from_index is None:
from_index = self.universe.start
if to_index is None:
to_index = self.universe.stop
if len(self._intervals) == 0:
return 0
if from_index >= to_index:
return 0
m = self._intervals.bisect_right(from_index)
n = self._intervals.bisect_right(to_index)
try:
from_index_index = self._intervals.index(from_index)
except ValueError:
from_index_index = None
try:
to_index_index = self._intervals.index(to_index)
except ValueError:
to_index_index = None
from_index_is_included = (
len(self._intervals) % 2 == 0 and m % 2 == 1 or len(self._intervals) % 2 == 1 and m % 2 == 0)
to_index_is_included = (
len(self._intervals) % 2 == 0 and n % 2 == 1 or len(self._intervals) % 2 == 1 and n % 2 == 0)
from_index_is_leftmost_included = from_index == 0 and from_index_is_included or from_index_index is not None and (
len(self._intervals) % 2 == 0 and from_index_index % 2 == 0
or len(self._intervals) % 2 == 1 and (from_index == 0 or from_index_index % 2 == 1))
to_index_right_of_excluded = to_index_index is not None and (
len(self._intervals) % 2 == 0 and to_index_index % 2 == 1
or len(self._intervals) % 2 == 1 and (to_index == 0 or to_index_index % 2 == 0))
if from_index_is_included:
if from_index_is_leftmost_included:
if to_index_is_included:
if m == 0:
to_remove = self._intervals[m:n]
endpoint = 0 if n == 0 else self._intervals[n - 1]
addendum = 0 if n == 0 else self._intervals[0]
self._revealed_count -= (to_index - endpoint) + addendum + sum(
b - a for a, b in zip(to_remove[1:-1:2], to_remove[2:-1:2]))
del self._intervals[m:n]
self._intervals.add(to_index)
else:
intermediates = self._intervals[m + 1:n - 1]
from_start, from_end = self._intervals[m - 1], self._intervals[m]
to_start, to_end = self._intervals[n - 1], self._intervals[n]
if m == n:
self._revealed_count -= to_index - from_start
self._intervals.remove(from_start)
self._intervals.add(to_index)
else:
self._revealed_count -= (from_end - from_start) + (to_index - self._intervals[n - 1]) + (
from_index - from_start) + sum(
b - a for a, b in zip(intermediates[::2], intermediates[1::2]))
del self._intervals[m + 1:n - 1] # intermediates
self._intervals.remove(from_start)
self._intervals.remove(from_end)
self._intervals.remove(to_start)
self._intervals.add(to_index)
else:
from_start = 0 if m == 0 else self._intervals[m - 1]
from_end = self._intervals[m]
self._revealed_count -= from_end - from_start
if from_start > 0:
self._intervals.remove(from_start)
self._intervals.remove(from_end)
else:
if to_index_is_included:
from_end = self._intervals[m]
to_start = self._intervals[n - 1]
if m == n:
self._revealed_count -= to_index - from_index
if from_index > 0:
self._intervals.add(from_index)
self._intervals.add(to_index)
else:
intermediates = self._intervals[m + 1:n - 1]
self._revealed_count -= (from_end - from_index) + (to_index - to_start) + sum(
b - a for a, b in zip(intermediates[::2], intermediates[1::2]))
del self._intervals[m + 1:n - 1] # intermediates
if from_index > 0:
self._intervals.add(from_index)
self._intervals.add(to_index)
self._intervals.remove(from_end)
self._intervals.remove(to_start)
else:
to_remove = self._intervals[m:n]
self._revealed_count -= self._intervals[m] - from_index + sum(b - a for a, b in zip(to_remove[1::2], to_remove[::2]))
del self._intervals[m:n]
if from_index != 0:
self._intervals.add(from_index)
else:
if to_index_is_included:
if to_index_right_of_excluded:
to_remove = self._intervals[m:n - 1]
del self._intervals[m:n - 1]
self._revealed_count -= sum(b - a for a, b in zip(to_remove[::2], to_remove[1::2]))
else:
to_remove = self._intervals[m:n]
del self._intervals[m:n]
self._intervals.add(to_index)
self._revealed_count -= (to_index - to_remove[0]) + sum(b - a for a, b in zip(to_remove[1::2], to_remove[::2]))
else:
to_remove = self._intervals[m:n]
del self._intervals[m:n]
self._revealed_count -= sum(b - a for a, b in zip(to_remove[::2], to_remove[1::2]))
return original_length - self._revealed_count
def exclude_virtual(self, from_index: Optional[int], to_index: Optional[int]):
if from_index is None or from_index < -len(self) or from_index >= len(self):
p_from_index = None
else:
p_from_index = self.virtual2physical(from_index)
if to_index is None or to_index < -len(self) or to_index >= len(self):
p_to_index = None
else:
p_to_index = self.virtual2physical(to_index)
return self.exclude(p_from_index, p_to_index)
def include(self, from_index: Optional[int], to_index: Optional[int]):
original_length = len(self)
if isinstance(from_index, int) and -self.universe.stop <= from_index < 0:
from_index = from_index % self.universe.stop
if isinstance(to_index, int):
if to_index > self.universe.stop:
return self.include(from_index, None)
if -self.universe.stop <= to_index < 0:
to_index = to_index % self.universe.stop
assert from_index is None or self.universe.start <= from_index <= self.universe.stop
assert to_index is None or self.universe.start <= to_index <= self.universe.stop
if from_index is None:
from_index = self.universe.start
if to_index is None:
to_index = self.universe.stop
if not self._intervals:
if from_index > 0:
self._intervals.add(from_index)
self._intervals.add(to_index)
self._revealed_count += to_index - from_index
return to_index - from_index
if from_index == to_index:
return 0
m = self._intervals.bisect_right(from_index)
n = self._intervals.bisect_right(to_index)
try:
from_index_index = self._intervals.index(from_index)
except ValueError:
from_index_index = None
from_index_is_included = (
len(self._intervals) % 2 == 0 and m % 2 == 1 or len(self._intervals) % 2 == 1 and m % 2 == 0)
to_index_is_included = (
len(self._intervals) % 2 == 0 and n % 2 == 1 or len(self._intervals) % 2 == 1 and n % 2 == 0)
from_index_right_of_included = from_index_index is not None and (
len(self._intervals) % 2 == 0 and from_index_index % 2 == 1
or len(self._intervals) % 2 == 1 and from_index_index % 2 == 0)
if from_index_is_included:
if to_index_is_included:
to_remove = self._intervals[m:n]
del self._intervals[m:n]
self._revealed_count += sum(b - a for a, b in zip(to_remove[::2], to_remove[1::2]))
else:
to_remove = self._intervals[m:n]
del self._intervals[m:n]
self._intervals.add(to_index)
self._revealed_count += (to_index - to_remove[-1]) + sum(b - a for a, b in zip(to_remove[1::2], to_remove[::2]))
else:
if to_index_is_included:
if from_index_right_of_included:
to_remove = self._intervals[m - 1:n]
del self._intervals[m - 1:n]
self._revealed_count += sum(b - a for a, b in zip(to_remove[::2], to_remove[1::2]))
else:
to_remove = self._intervals[m:n]
del self._intervals[m:n]
self._intervals.add(from_index)
self._revealed_count += (to_remove[0] - from_index) + sum(b - a for a, b in zip(to_remove[1::2], to_remove[::2]))
else:
if from_index_right_of_included:
intermediates = self._intervals[m:n]
del self._intervals[m:n] # intermediates
self._intervals.remove(from_index)
self._intervals.add(to_index)
self._revealed_count += (to_index - from_index) - sum(b - a for a, b in zip(intermediates[::2], intermediates[1::2]))
else:
to_remove = self._intervals[m:n]
del self._intervals[m:n]
if from_index > 0:
self._intervals.add(from_index)
self._intervals.add(to_index)
self._revealed_count += (to_index - from_index) - sum(b - a for a, b in zip(to_remove[::2], to_remove[1::2]))
return len(self) - original_length
def include_partially(self, from_index: Optional[int], to_index: Optional[int], count: Union[int, tuple]):
if isinstance(count, int):
return self.include_partially(from_index, to_index, (count, count))
head_count, tail_count = count
head_revealed_count = self._include_partially_from_left(from_index, to_index, head_count)
tail_revealed_count = self._include_partially_from_right(from_index, to_index, tail_count)
return head_revealed_count + tail_revealed_count
def _include_partially_from_left(self, from_index: int, to_index: int, count: int):
if count == 0:
return 0
from_index, to_index = self._normalized_range(from_index, to_index)
subsel = self._spanning_subslice(from_index, to_index).complement().subslice(from_index, to_index)
revealed_count = 0
for covered_start, covered_stop in subsel.pairs():
coverage = covered_stop - covered_start
if revealed_count + coverage < count:
self.include(covered_start, covered_stop)
revealed_count += coverage
else:
self.include(covered_start, covered_start + count - revealed_count)
revealed_count = count
break
return revealed_count
def _include_partially_from_right(self, from_index: int, to_index: int, count: int):
if count == 0:
return 0
from_index, to_index = self._normalized_range(from_index, to_index)
subsel = self._spanning_subslice(from_index, to_index).complement().subslice(from_index, to_index)
revealed_count = 0
for covered_start, covered_stop in reversed(list(subsel.pairs())):
coverage = covered_stop - covered_start
if revealed_count + coverage < count:
self.include(covered_start, covered_stop)
revealed_count += coverage
else:
self.include(covered_stop - (count - revealed_count), covered_stop)
revealed_count = count
break
return revealed_count
def include_expand(self, from_index: Optional[int], to_index: Optional[int], count: Union[int, Tuple[int, int]]):
if isinstance(count, int):
return self.include_expand(from_index, to_index, (count, count))
if count == (0, 0):
return 0
head_count, tail_count = count
revealed_counter = 0
gaps = self.complement().subslice(from_index, to_index)
for a, b in gaps.pairs():
if b < self.universe.stop:
revealed_counter += self._include_partially_from_right(a, b, head_count)
if a > self.universe.start:
revealed_counter += self._include_partially_from_left(a, b, tail_count)
return revealed_counter
def _previous_slice(self, sl: slice):
""" :return The revealed or covered slice immediately to the left of @sl.
:raise ValueError if there is none. """
if sl.start == self.universe.start:
raise ValueError("There is no slice to the left of {}.".format(sl))
# TODO O(n) -> O(1)
zero_or_one = [s for s in self._intervals + self.complement()._intervals if s.stop == sl.start]
if len(zero_or_one) == 1:
return zero_or_one[0]
else:
raise ValueError("Slice not found: {}.".format(sl))
def _next_slice(self, sl: slice):
""" :return The revealed or covered slice immediately to the right of @sl.
:raise ValueError if there is none. """
if sl.stop == self.universe.stop:
raise ValueError("There is no slice to the right of {}.".format(sl))
# TODO O(n)
zero_or_one = [s for s in self._intervals + self.complement()._intervals if s.start == sl.stop]
if len(zero_or_one) == 1:
return zero_or_one[0]
else:
raise ValueError("Slice not found: {}.".format(sl))
def include_virtual(self, from_index, to_index):
if from_index is None or from_index < -len(self) or from_index >= len(self):
p_from_index = None
else:
p_from_index = self.virtual2physical(from_index)
if to_index is None or to_index < -len(self) or to_index >= len(self):
p_to_index = None
else:
p_to_index = self.virtual2physical(to_index)
return self.include(p_from_index, p_to_index)
def include_partially_virtual(self, from_index: Optional[int], to_index: Optional[int], count: Union[int, tuple]):
if from_index is None or from_index < -len(self) or from_index >= len(self):
p_from_index = None
else:
p_from_index = self.virtual2physical(from_index)
if to_index is None or to_index < -len(self) or to_index >= len(self):
p_to_index = None
else:
p_to_index = self.virtual2physical(to_index)
return self.include_partially(p_from_index, p_to_index, count)
# FIXME Inconsistent with reversed(selection). Should probably make this use the default implementation and instead
# rewrite this one to iter_slices or something.
def __iter__(self):
for a, b in self.pairs():
yield a, b # FIXME should probably generate slices instead, or every index
def complement(self):
if len(self._intervals) >= 1 and self._intervals[-1] == self.universe.stop:
return Selection(universe=self.universe, intervals=self._intervals[:-1],
_length=self.universe.stop - len(self))
return Selection(universe=self.universe, intervals=self._intervals.union([self.universe.stop]),
_length=self.universe.stop - len(self))
def _normalized_range(self, from_index: Optional[int], to_index: Optional[int]) -> Tuple[int, int]:
""" For any range [@from_index, @to_index) where the indices are either None or any integer, returns the
equivalent range [x, y) such that either 0 <= x < y <= upper_bound or x = y = 0. The ranges are equivalent in
the sense that when using them to slice this selection, they produce the same sub-selection. """
if from_index is None or from_index <= -self.universe.stop:
from_index = self.universe.start
elif from_index > self.universe.stop:
from_index = self.universe.stop
elif -self.universe.stop <= from_index < 0:
from_index = self.universe.stop - from_index
if to_index is None or to_index >= self.universe.stop:
to_index = self.universe.stop
elif -self.universe.stop <= to_index < 0:
to_index = self.universe.stop - to_index
elif to_index < -self.universe.stop:
to_index = self.universe.start
if from_index >= to_index:
from_index, to_index = (0, 0)
return from_index, to_index
def subslice(self, from_index: Optional[int], to_index: Optional[int]):
from_index, to_index = self._normalized_range(from_index, to_index)
sel = self._spanning_subslice(from_index, to_index)
if len(sel._intervals) % 2 == 0:
if len(sel) > 0:
if sel._intervals[0] < from_index < sel._intervals[1]:
sel._revealed_count -= from_index - sel._intervals[0]
del sel._intervals[0]
sel._intervals.add(from_index)
if sel._intervals[-2] < to_index < sel._intervals[-1]:
sel._revealed_count -= sel._intervals[-1] - to_index
del sel._intervals[-1]
sel._intervals.add(to_index)
else:
if 0 < from_index < sel._intervals[0]:
sel._revealed_count -= from_index
sel._intervals.add(from_index)
if (len(sel._intervals) == 1 and to_index < sel._intervals[-1]
or len(sel._intervals) >= 2 and sel._intervals[-2] < to_index < sel._intervals[-1]):
sel._revealed_count -= sel._intervals[-1] - to_index
del sel._intervals[-1]
sel._intervals.add(to_index)
return sel
def _spanning_subslice(self, from_index: int, to_index: int):
""" :return A Selection whose set of revealed slices is a subset of that of this Selection such that every index
in [from_index, to_index) is either on some slice in the subset, or on a gap. """
if from_index >= to_index:
return Selection(universe=deepcopy(self.universe), intervals=[])
m = self._intervals.bisect_right(from_index)
if len(self._intervals) % 2 == 0:
n = self._intervals.bisect_left(to_index)
intervals = self._intervals[m - (m % 2):n + (n % 2)]
else:
n = self._intervals.bisect_right(to_index)
a = max(0, m - ((m + 1) % 2))
b = n + ((n + 1) % 2)
intervals = self._intervals[a:b]
sel = Selection(universe=deepcopy(self.universe), intervals=intervals)
return sel
def _slow_subslice(self, from_index: Optional[int], to_index: Optional[int]):
sel = self.deepcopy()
if isinstance(from_index, int):
sel.exclude(None, from_index)
if isinstance(to_index, int):
sel.exclude(to_index, None)
return sel
def _interval_index(self, pindex):
""" :return n if the nth interval edge is the smallest number such that @pindex < n (zero-indexed). """
lower = 0
upper = len(self._intervals) - 1
while lower <= upper:
middle = (lower + upper) // 2
midsl = self._intervals[middle]
if pindex < midsl.start:
upper = middle - 1
elif midsl.stop <= pindex:
lower = middle + 1
else: # midsl.start <= pindex < midsl.stop:
return middle
raise IndexError("{} is not in any interval.".format(pindex))
def select(self, listlike):
# TODO only works for stringlike objects
lst = []
for interval in self.slices():
lst.append(listlike[interval])
selection = listlike[0:0].join(lst)
return selection
def physical2virtual(self, pindex: int):
vindex = 0
for a, b in self.pairs():
if a <= pindex < b:
vindex += pindex - a
return vindex
vindex += b - a
raise IndexError("Physical index {} out of bounds for selection {}".format(pindex, self))
# TODO: O(n) -> O(log(n)) (using another sorted set for cumulative lengths?)
def virtual2physical(self, vindex: int): # TODO -> virtualint2physical
""" :return the integer n such that where the @vindex'th revealed element is the nth element. If
@vindex < 0, @vindex is interpreted as (number of revealed elements) + @vindex.
"""
if vindex < -len(self):
raise IndexError(
"Got index {}, expected it to be within range [{},{})".format(vindex, -len(self), len(self)))
elif vindex < 0:
return self.virtual2physical(len(self) + vindex)
cumlength = 0
for a, b in self.pairs():
cumlength += b - a
if vindex < cumlength:
pindex = b - (cumlength - vindex)
if a <= pindex < b:
return pindex
else:
break
raise IndexError("Virtual index {} out of bounds for selection {}".format(vindex, self))
def virtual2physicalselection(self, vslice: slice) -> 'Selection': # TODO -> virtualslice2physical
""" :return the sub-Selection that is the intersection of this selection and @vslice. """
if not self._intervals or vslice.stop == 0:
return Selection(self.universe, revealed=[])
if vslice.start is None:
a = self.virtual2physical(0)
elif -len(self) <= vslice.start < len(self):
a = self.virtual2physical(vslice.start)
elif vslice.start >= len(self):
a = self._intervals[-1]
else:
raise ValueError("Unexpected slice start: {}".format(vslice))
if vslice.stop is None or vslice.stop >= len(self):
b = self._intervals[-1] - 1
elif -len(self) <= vslice.stop < len(self):
b = self.virtual2physical(vslice.stop - 1)
else:
raise ValueError("Unexpected slice stop: {}".format(vslice))
# INV: a is the physical index of the first element, b is the physical index of the last element
if b < a:
return Selection(universe=self.universe, revealed=[])
m = self._intervals.bisect_right(a)
n = self._intervals.bisect_right(b)
intervals = SortedSet([a] + self._intervals[m:n] + [b + 1])
return Selection(universe=self.universe, intervals=intervals)
def virtualselection2physical(self, vselection: 'Selection'): # TODO -> virtualslice2physical
""" :return the sub-Selection that is the intersection of this selection and @vselection. """
intervals = []
for start, stop in vselection:
for a, b in self.virtual2physicalselection(slice(start, stop)):
intervals.append(slice(a, b))
return Selection(universe=self.universe, revealed=intervals)
def stretched(self, from_index: Optional[int], to_index: Optional[int]): # TODO remove?
""" :return A potentially shrinked deep copy of this selection, delimited by the universe
[@from_index, @to_index). """
m = self._intervals.bisect_right(from_index)
n = self._intervals.bisect_right(to_index)
intervals = self._intervals[m:n]
return Selection(universe=slice(from_index, to_index), intervals=intervals)
def __getitem__(self, item):
return self.virtual2physical(item)
@staticmethod
def _compute_len(sortedset: SortedSet):
""" :return The sum of the lengths of every slice in @slicelist. """
if len(sortedset) == 0:
return 0
elif len(sortedset) % 2 == 0:
return sum(sortedset[i + 1] - sortedset[i] for i in range(0, len(sortedset), 2))
return sortedset[0] + sum(sortedset[i + 1] - sortedset[i] for i in range(1, len(sortedset), 2))
def __len__(self):
return self._revealed_count
def __eq__(self, other):
return repr(self) == repr(other)
def __mul__(self, other: int):
if other == 0:
return Selection(universe=slice(0, 0), revealed=[])
scaled_universe = slice(self.universe.start * other, self.universe.stop * other)
scaled_revealed = [other * x for x in self._intervals]
return Selection(universe=scaled_universe, intervals=scaled_revealed)
def __rmul__(self, other):
return self.__mul__(other)
def __repr__(self):
return "{}(universe={}, intervals={})".format(self.__class__.__name__, self.universe, self._intervals)
def __str__(self):
return repr(self)
def deepcopy(self):
""" :return A deep copy of this object. """
return Selection(universe=deepcopy(self.universe), intervals=deepcopy(self._intervals))
def test5():
"""
有序的集合：SortedSet
网址：http://www.grantjenks.com/docs/sortedcontainers/sortedset.html
"""
from sortedcontainers import SortedSet
# 创建 SortedSet
ss = SortedSet([3, 1, 2, 5, 4])
print(ss) # SortedSet([1, 2, 3, 4, 5])
from operator import neg
ss1 = SortedSet([3, 1, 2, 5, 4], neg)
print(ss1) # SortedSet([5, 4, 3, 2, 1], key=<built-in function neg>)
# SortedSet 转为 list/tuple/set
print(list(ss)) # SortedSet转为list [1, 2, 3, 4, 5]
print(tuple(ss)) # SortedSet转为tuple (1, 2, 3, 4, 5)
print(set(ss)) # SortedSet转为set {1, 2, 3, 4, 5}
# 插入、删除元素
ss.discard(-1) # 删除不存在的元素不报错
ss.remove(1) # 删除不存在的元素报错, KeyError
ss.discard(3) # SortedSet([1, 2, 4, 5])
ss.add(-10) # SortedSet([-10, 1, 2, 4, 5])
# 返回第一个和最后一个元素
print(ss[0]) # -10
print(ss[-1]) # 5
# 遍历 set
for e in ss:
print(e, end=", ") # -10, 2, 4, 5,
print()
# set 中判断某元素是否存在
print(2 in ss) # True
# bisect_left() / bisect_right()
print(ss.bisect_left(4)) # 返回大于等于4的最小元素对应的下标 2
print(ss.bisect_right(4)) # 返回大于4的最小元素对应的下标 3
# 清空 set
ss.clear()
print(len(ss)) # 0
print(len(ss) == 0) # True
"""
无序的集合: set
"""
# 集合的定义：集合是不可变的，因此集合中元素不能是list
A = {"hi", 2, ("we", 24)}
B = set() # 空集合的定义，不能使用B = {}定义集合，这样是字典的定义
# 集合间的操作, 下面的运算法符都可以写成 op= 的形式
print("---------------------------------------")
S = {1, 2, 3}
T = {3, 4, 5}
print(S & T) # 交集，返回一个新集合，包括同时在集合S和T中的元素
print(S | T) # 并集，返回一个新集合，包括在集合S和T中的所有元素
print(S - T) # 差集，返回一个新集合，包括在集合S但不在T中的元素
print(S ^ T) # 补集，返回一个新集合，包括集合S和T中的非相同元素
# 集合的包含关系
print("---------------------------------------")
C = {1, 2}
D = {1, 2}
print(C <= D) # C是否是D的子集 True
print(C < D) # C是否是D的真子集 False
print(C >= D) # D是否是C的子集 True
print(C > D) # D是否是C的真子集 False
# 集合的处理方法
print("---------------------------------------")
S = {1, 2, 3, 5, 6}
S.add(4) # 如果x不在集合S中，将x增加到S
S.discard(1) # 移除S中元素x，如果x不在集合S中，不报错
S.remove(2) # 移除S中元素x，如果x不在集合S中，产生KeyError异常
for e in S: # 遍历
print(e, end=",")
print()
print(S.pop()) # 从S中随机弹出一个元素，S长度减1，若S为空产生KeyError异常
print(S.copy()) # 返回集合S的一个副本, 对该副本的操作不会影响S
print(len(S)) # 返回集合S的元素个数
print(5 in S) # 判断S中元素x， x在集合S中，返回True，否则返回False
print(5 not in S) # 判断S中元素x， x在集合S中，返回True，否则返回False
S.clear() # 移除S中所有元素
# 작지만 큰 숫자
## n개의 숫자로 이루어진 수열이 하나 주어지고, 그 이후 m개의 질의가 주어집니다
## 각 질의마다 하나의 숫자가 주어진다 했을 때, 순서대로 수열 내에서 주어진 숫자보다 같거나 작은 숫자들 중 최댓값을 하나 골라 제거하는 것을 반복하는 프로그램을 작성해보세요
## 단, 같거나 작은 숫자가 없는 경우에는 제거하지 않고 넘어갑니다.
from sortedcontainers import SortedSet
s = SortedSet() # treeset
numSize = list(map(int, input().split()))
numList = list(map(int, input().split()))
numSet = SortedSet(numList)
cmd = list(map(int, input().split()))
for i in cmd:
if i in numSet:
print(i)
numSet.remove(i)
else:
idx = numSet.bisect_right(i) - 1
if idx < 0:
print(-1)
else:
print(numSet[idx])
numSet.remove(numSet[idx])
s = SortedSet() # treeset
cmd = list(map(int,input().split()))
peopleWant = list(map(int, input().split()))
peopleWant.sort(reverse=True)
chairSet = SortedSet()
chair = []
for i in range(cmd[1]):
chairSet.add(i+1)
resultCount = 0
for i in peopleWant:
idx = chairSet.bisect_right(i)-1
if i in chairSet:
chairSet.remove(i)
resultCount += 1
else:
if idx < 0:
#print(i)
break
else:
chairSet.remove(chairSet[idx])
resultCount += 1
print(resultCount)
class MutableIntervalDict(
# pylint: disable=unsubscriptable-object
IntervalDict[atomic.TO, V],
MutableMapping[atomic.Interval[atomic.TO], V],
):
"""
Mutable Interval Dictionary class.
The :class:`MutableIntervalDict` class (which inherits from the
:class:`IntervalDict` class) is designed to hold mutable dict of disjoint sorted
intervals.
Note
----
Les :math:`n` (or :math:`n_0`) the number of intervals of the *self* variable
and :math:`m` the number of intervals in the *other* variable. Let
:math:`n_1, ... n_k` the number of intervals for methods with multiple
arguments.
The complexity in time of methods is:
========================= ====================================================
Methods Average case
========================= ====================================================
:meth:`__setitem__` :math:`O(n)`
:meth:`__delitem__` :math:`O(n)`
:meth:`__ior__` :math:`O(m\\log(n+m))`
:meth:`update` :math:`O((\\sum_{i=1}^kn_i)\\log(\\sum_{i=0}^kn_i))`
:meth:`clear` :math:`O(1)`
========================= ====================================================
"""
__slots__ = ("_default", "_operator", "_strict")
def __init__(
self,
iterable: Optional[Union[IntervalDict[atomic.TO, V], Mapping[
atomic.IntervalValue[atomic.TO],
V], Iterable[Tuple[atomic.IntervalValue[atomic.TO], V]], ]] = None,
default: Optional[Callable[[], V]] = None,
operator: Optional[Callable[[V, V], V]] = None,
strict: Optional[bool] = True,
) -> None:
"""
Initialize a :class:`MutableIntervalDict` instance.
Arguments
---------
iterable: :class:`Iterable <python:typing.Iterable>`
An optional iterable that can be converted to a dictionary of (
interval, value).
Keyword arguments
-----------------
default: :class:`Callable[[], V] <python:typing.Callable>`, optional
The default factory.
operator: :class:`Callable[[V, V], V] <python:typing.Callable>`
The operator function.
strict: bool
:data:`False <python:False>` if ``operator`` is a commutative and
associative law on ``V``.
Note
----
If ``operator`` is a commutative and associative law on ``V``,
the complexity in time is much faster if ``strict`` is set to
:data:`False <python:False>`.
Examples
--------
>>> from part import MutableIntervalDict
>>> a = MutableIntervalDict[int, set](
... operator=lambda x, y: x | y,
... strict=False
... )
>>> a.update({(1, 10): {1}})
>>> print(a)
{'[1;10)': {1}}
>>> a.update({(5, 20): {2}})
>>> print(a)
{'[1;5)': {1}, '[5;10)': {1, 2}, '[10;20)': {2}}
>>> a.update({(10, 30): {1}})
>>> print(a)
{'[1;5)': {1}, '[5;10)': {1, 2}, '[10;20)': {1, 2}, '[20;30)': {1}}
>>> print(a.compress())
{'[1;5)': {1}, '[5;20)': {1, 2}, '[20;30)': {1}}
"""
super().__init__()
self._default = default
self._operator = operator
self._strict = strict
if isinstance(iterable, IntervalDict):
self._mapping = iterable._mapping.copy()
self._intervals = SortedSet(iterable._intervals)
else:
self._mapping = {}
self._intervals: SortedSet = SortedSet()
if isinstance(iterable, dict):
self.update(*({key: value} for key, value in iterable.items()))
elif iterable is not None:
self.update(*({
key: value
} for key, value in iterable)) # type: ignore
def __getitem__(self, key: Union[slice,
atomic.IntervalValue[atomic.TO]]) -> V:
"""
Return a value using either a slice or an interval value.
Arguments
---------
key: Union[IntervalValue, slice]
The interval requested.
Returns
-------
The found value
Raises
------
KeyError
If the *key* is out of range.
"""
try:
return super().__getitem__(key)
except KeyError:
if self._default is not None:
value = self._default()
self[key] = value
return value
raise
def __setitem__(self, key: Union[slice, atomic.IntervalValue[atomic.TO]],
value: V) -> None:
"""
Set a value using either a slice or an interval value.
Arguments
---------
key: Union[IntervalValue, slice]
The interval requested.
Raises
------
KeyError
If the *key* is out of range.
Examples
--------
>>> from part import MutableIntervalDict
>>> a = MutableIntervalDict[int, int](
... {(10, 15): 1, (20, 25): 2, (30, 35): 3}
... )
>>> print(a)
{'[10;15)': 1, '[20;25)': 2, '[30;35)': 3}
>>> a[12] = 4
>>> print(a)
{'[10;12)': 1, '[12;12]': 4, '(12;15)': 1, '[20;25)': 2, '[30;35)': 3}
>>> a[13:31] = 5
>>> print(a)
{'[10;12)': 1, '[12;12]': 4, '(12;13)': 1, '[13;31)': 5, '[31;35)': 3}
>>> a[:] = 0
>>> print(a)
{'(-inf;+inf)': 0}
"""
interval = self._remove(key)
if interval:
self._intervals.add(interval)
self._mapping[interval] = value
def __delitem__(
self, key: Union[slice, atomic.IntervalValue[atomic.TO]]) -> None:
"""
Delete a value using either a slice or an interval value.
Arguments
---------
key: Union[IntervalValue, slice]
The interval requested.
Raises
------
KeyError
If the *key* is out of range.
Examples
--------
>>> from part import MutableIntervalDict
>>> a = MutableIntervalDict[int, int](
... {(10, 15): 1, (20, 25): 2, (30, 35): 3}
... )
>>> print(a)
{'[10;15)': 1, '[20;25)': 2, '[30;35)': 3}
>>> del a[12]
>>> print(a)
{'[10;12)': 1, '(12;15)': 1, '[20;25)': 2, '[30;35)': 3}
>>> del a[13:31]
>>> print(a)
{'[10;12)': 1, '(12;13)': 1, '[31;35)': 3}
>>> del a[:]
>>> print(a)
{}
"""
self._remove(key)
def _remove(self, key):
# pylint: disable=protected-access
interval = IntervalDict._interval(key)
if interval:
start = self._start(interval)
stop = self._stop(interval)
for index in range(start, stop):
del self._mapping[self._intervals[index]]
del self._intervals[start:stop]
return interval
def _add(self, interval, value):
if self._operator is None:
self[interval] = value
else:
intervals = list(self.select(interval, strict=False))
for another in ((interval & found)[0] for found in intervals):
self[another] = self._operator(self[another], value)
for another in sets.FrozenIntervalSet[atomic.TO](
[interval]) - sets.FrozenIntervalSet[atomic.TO](intervals):
self[another] = value
def __or__(self, other) -> "MutableIntervalDict[atomic.TO, V]":
"""
Construct a new dictionary using self and the *other*.
Arguments
---------
other: :class:`IntervalDict`
Another interval dict.
Returns
-------
:class:`MutableIntervalDict`
The new :class:`IntervalDict`.
Examples
--------
>>> from part import MutableIntervalDict
>>> a = MutableIntervalDict[int, int](
... {(10, 15): 1, (20, 25): 2, (30, 35): 3},
... operator=lambda x, y: x + y,
... strict=False
... )
>>> print(a | FrozenIntervalDict[int, int]({(15, 22): 4}))
{'[10;15)': 1, '[15;20)': 4, '[20;22)': 6, '[22;25)': 2, '[30;35)': 3}
"""
if not isinstance(other, IntervalDict):
return NotImplemented
result = self.__class__(self,
default=self._default,
operator=self._operator,
strict=self._strict)
result.update(other)
return result
def __ior__(self, other) -> "MutableIntervalDict[atomic.TO, V]":
"""
Update self with the *other*.
Arguments
---------
other: :class:`IntervalDict`
Another interval dict.
Returns
-------
:class:`MutableIntervalDict`
The updated :class:`MutableIntervalDict`.
Examples
--------
>>> from part import MutableIntervalDict
>>> a = MutableIntervalDict[int, int](
... operator=lambda x, y: x + y,
... strict=False
... )
>>> a |= MutableIntervalDict[int, int]({(1, 10): 1})
>>> print(a)
{'[1;10)': 1}
>>> a |= MutableIntervalDict[int, int]({(5, 20): 2})
>>> print(a)
{'[1;5)': 1, '[5;10)': 3, '[10;20)': 2}
>>> a |= MutableIntervalDict[int, int]({(10, 30): 3})
>>> print(a)
{'[1;5)': 1, '[5;10)': 3, '[10;20)': 5, '[20;30)': 3}
"""
if not isinstance(other, IntervalDict):
return NotImplemented
self.update(other)
return self
# pylint: disable=arguments-differ,signature-differs
def update( # type: ignore
self,
*args: Union[IntervalDict, Mapping[atomic.IntervalValue[atomic.TO], V],
Iterable[Tuple[atomic.IntervalValue[atomic.TO], V]], ],
) -> None:
"""
Update the dict.
Arguments
---------
*args: :class:`Iterable <python:typing.Iterable>`
An iterable of :class:`IntervalDict` or valid iterable for an interval
dictionary creation.
Raises
------
TypeError
if an argument is not iterable.
Note
----
If the parameter ``strict`` used in the constructor is :data:`False
<python:False>`, the complexity is in :math:`O(n\\,log(n)\\,k\\,\\lambda)`
where:
* :math:`n` is the length of ``*args``;
* :math:`k` is the number of output intervals;
* :math:`\\lambda` is the the cost of the ``operator`` parameter used in
the constructor.
Examples
--------
>>> from part import MutableIntervalDict
>>> from operator import add
>>> a = MutableIntervalDict[int, int](
... operator=add,
... default=lambda: 0,
... )
>>> a.update({(1, 10): 1})
>>> print(a)
{'[1;10)': 1}
>>> a.update(
... FrozenIntervalDict[int, int]({(5, 20): 2}),
... FrozenIntervalDict[int, int]({(10, 30): 3})
... )
>>> print(a)
{'[1;5)': 1, '[5;10)': 3, '[10;20)': 5, '[20;30)': 3}
>>> a = MutableIntervalDict[int, set](
... operator=lambda x, y: x | y,
... strict=False
... )
>>> a.update({(1, 10): {1}})
>>> print(a)
{'[1;10)': {1}}
>>> a.update(
... FrozenIntervalDict[int, set]({(5, 20): {2}}),
... FrozenIntervalDict[int, set]({(10, 30): {3}})
... )
>>> print(a)
{'[1;5)': {1}, '[5;10)': {1, 2}, '[10;20)': {2, 3}, '[20;30)': {3}}
"""
# TODO determine complexity
strict = self._strict
operator = self._operator
if strict or operator is None:
self._strict_update(*args)
else:
self._enhanced_update(*args)
def _strict_update(
self,
*args: Union[IntervalDict, Mapping[atomic.IntervalValue[atomic.TO], V],
Iterable[Tuple[atomic.IntervalValue[atomic.TO], V]], ],
) -> None:
for other in args:
if isinstance(other, collections.abc.Mapping):
for key, value in other.items():
self._add(atomic.Atomic.from_value(key), value)
elif isinstance(other, collections.abc.Iterable):
for key, value in other: # type: ignore
self._add(atomic.Atomic.from_value(key), value)
else:
raise TypeError(f"{type(other)} object is not iterable")
# pylint: disable=protected-access
@classmethod
def _rest(cls, rest, cursors, index, element):
if cursors[index] < len(element):
interval = element._intervals[cursors[index]]
value = element._mapping[interval]
rest.add((interval.lower, interval.upper, index, value))
# pylint: disable=protected-access
def _create(self, *args):
# Create a list of non empty IntervalDict
elements = []
for element in itertools.chain([self], args):
if not isinstance(element, IntervalDict):
element = FrozenIntervalDict(element)
if element:
elements.append(element)
cursors = [0] * len(elements)
current = SortedList()
rest = SortedList()
for index, element in enumerate(elements):
self._rest(rest, cursors, index, element)
return (elements, cursors, current, rest)
# pylint: disable=too-many-arguments,protected-access
@classmethod
def _next(cls, upper, elements, cursors, current, rest):
# Remove useless elements from current
while current and current[0][0] == upper:
(_, _, index, _) = current[0]
del current[0]
cursors[index] += 1
element = elements[index]
cls._rest(rest, cursors, index, element)
if current or not rest:
lower = upper.next()
else:
lower = rest[0][0]
# Move elements from rest to current
while rest and rest[0][0] == lower:
(lower, upper, index, value) = rest[0]
del rest[0]
current.add((upper, lower, index, value))
if current:
upper = current[0][0]
if rest:
upper = min(upper, rest[0][0].prev())
return (lower, upper)
def _enhanced_update(
self,
*args: Union[IntervalDict, Mapping[atomic.IntervalValue[atomic.TO], V],
Iterable[Tuple[atomic.IntervalValue[atomic.TO], V]], ],
) -> None:
intervals = []
mapping = {}
(elements, cursors, current, rest) = self._create(*args)
(lower, upper) = self._next(-atomic.INFINITY, elements, cursors,
current, rest)
while current:
interval = atomic.Interval[atomic.TO](
lower_value=lower.value,
lower_closed=lower.type == 0,
upper_value=upper.value,
upper_closed=upper.type == 0,
)
value = reduce(
self._operator,
(value for (_, _, _, value) in current) # type: ignore
)
intervals.append(interval)
mapping[interval] = value
(lower, upper) = self._next(upper, elements, cursors, current,
rest)
self._intervals = SortedSet(intervals)
self._mapping = mapping
def clear(self) -> None:
"""Remove all items from self (same as del self[:])."""
self._intervals = SortedSet()
self._mapping = {}
def _start(self, interval):
start = self._intervals.bisect_left(interval)
if start < len(self._intervals):
start_interval = self._intervals[start]
start_value = self._mapping[start_interval]
if self._intervals[start].overlaps(interval, strict=False):
del self._intervals[start]
del self._mapping[start_interval]
if self._insert(
(
start_interval.lower_value,
interval.lower_value,
start_interval.lower_closed,
not interval.lower_closed,
),
start_value,
):
start += 1
elif interval.overlaps(self._intervals[start], strict=False):
del self._intervals[start]
del self._mapping[start_interval]
if self._insert(
(
interval.upper_value,
start_interval.upper_value,
not interval.upper_closed,
start_interval.upper_closed,
),
start_value,
):
start += 1
elif interval.during(self._intervals[start], strict=False):
del self._intervals[start]
del self._mapping[start_interval]
if self._insert(
(
start_interval.lower_value,
interval.lower_value,
start_interval.lower_closed,
not interval.lower_closed,
),
start_value,
):
start += 1
if self._insert(
(
interval.upper_value,
start_interval.upper_value,
not interval.upper_closed,
start_interval.upper_closed,
),
start_value,
):
start += 1
return start
def _stop(self, interval):
stop = self._intervals.bisect_right(interval)
if 0 < stop <= len(self._intervals):
stop -= 1
stop_interval = self._intervals[stop]
stop_value = self._mapping[stop_interval]
if self._intervals[stop].during(interval, strict=False):
del self._intervals[stop]
del self._mapping[stop_interval]
elif interval.overlaps(self._intervals[stop], strict=False):
del self._intervals[stop]
del self._mapping[stop_interval]
self._insert(
(
interval.upper_value,
stop_interval.upper_value,
not interval.upper_closed,
stop_interval.upper_closed,
),
stop_value,
)
return stop
def _insert(self, key, value):
interval = atomic.Atomic.from_tuple(key)
if interval:
self._intervals.add(interval)
self._mapping[interval] = value
return True
return False
def _bisect_left(self, search) -> int:
return self._intervals.bisect_left(search)
def _update_intervals(self, intervals) -> None:
self._intervals = SortedSet(intervals)
"""
https://leetcode.com/discuss/general-discussion/452863/how-do-you-deal-with-no-treeset-or-treemap-in-python
Python3 可以使用 sortedcontainers 库来实现 Java 中 TreeSet, TreeMap, Collections.sort(list)
"""
from sortedcontainers import SortedList
sl = SortedList(['e', 'a', 'c', 'd', 'b'])
print(sl) # SortedList(['a', 'b', 'c', 'd', 'e'])
sl *= 100
print(sl.count('c')) # 100
print(sl[-3:]) # ['e', 'e', 'e']
from sortedcontainers import SortedDict
sd = SortedDict({'c': 3, 'a': 1, 'b': 2})
print(sd) # SortedDict({'a': 1, 'b': 2, 'c': 3})
print(sd.popitem(index=-1)) # ('c', 3)
from sortedcontainers import SortedSet
ss = SortedSet('abracadabra')
print(ss) # SortedSet(['a', 'b', 'c', 'd', 'r'])
print(ss.bisect_left('c')) # 2
print(ss.bisect_right('c')) # 3
print(ss.bisect_left('f')) # 4
print(ss.bisect_right('f')) # 4
