Python/SortedMultiSet.py at main · achapi/Python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
importmath
frombisectimportbisect_left, bisect_right, insort
fromtypingimportGeneric, Iterable, Iterator, TypeVar, Union, List
T=TypeVar('T')

classSortedMultiset(Generic[T]):
BUCKET_RATIO=50
REBUILD_RATIO=170

def_build(self, a=None) ->None:
"Evenly divide `a` into buckets."
ifaisNone: a=list(self)
size=self.size=len(a)
bucket_size=int(math.ceil(math.sqrt(size/self.BUCKET_RATIO)))
self.a= [a[size*i//bucket_size : size* (i+1) //bucket_size] foriinrange(bucket_size)]

def__init__(self, a: Iterable[T] = []) ->None:
"Make a new SortedMultiset from iterable. / O(N) if sorted / O(N log N)"
a=list(a)
ifnotall(a[i] <=a[i+1] foriinrange(len(a) -1)):
a=sorted(a)
self._build(a)

def__iter__(self) ->Iterator[T]:
foriinself.a:
forjini: yieldj

def__reversed__(self) ->Iterator[T]:
foriinreversed(self.a):
forjinreversed(i): yieldj

def__len__(self) ->int:
returnself.size

def__repr__(self) ->str:
return"SortedMultiset"+str(self.a)

def__str__(self) ->str:
s=str(list(self))
return"{"+s[1 : len(s) -1] +"}"

def_find_bucket(self, x: T) ->List[T]:
"Find the bucket which should contain x. self must not be empty."
forainself.a:
ifx<=a[-1]: returna
returna

def__contains__(self, x: T) ->bool:
ifself.size==0: returnFalse
a=self._find_bucket(x)
i=bisect_left(a, x)
returni!=len(a) anda[i] ==x

defcount(self, x: T) ->int:
"Count the number of x."
returnself.index_right(x) -self.index(x)

defadd(self, x: T) ->None:
"Add an element. / O(√N)"
ifself.size==0:
self.a= [[x]]
self.size=1
return
a=self._find_bucket(x)
insort(a, x)
self.size+=1
iflen(a) >len(self.a) *self.REBUILD_RATIO:
self._build()

defdiscard(self, x: T) ->bool:
"Remove an element and return True if removed. / O(√N)"
ifself.size==0: returnFalse
a=self._find_bucket(x)
i=bisect_left(a, x)
ifi==len(a) ora[i] !=x: returnFalse
a.pop(i)
self.size-=1
iflen(a) ==0: self._build()
returnTrue

deflt(self, x: T) ->Union[T, None]:
"Find the largest element < x, or None if it doesn't exist."
forainreversed(self.a):
ifa[0] <x:
returna[bisect_left(a, x) -1]

defle(self, x: T) ->Union[T, None]:
"Find the largest element <= x, or None if it doesn't exist."
forainreversed(self.a):
ifa[0] <=x:
returna[bisect_right(a, x) -1]

defgt(self, x: T) ->Union[T, None]:
"Find the smallest element > x, or None if it doesn't exist."
forainself.a:
ifa[-1] >x:
returna[bisect_right(a, x)]

defge(self, x: T) ->Union[T, None]:
"Find the smallest element >= x, or None if it doesn't exist."
forainself.a:
ifa[-1] >=x:
returna[bisect_left(a, x)]

def__getitem__(self, x: int) ->T:
"Return the x-th element, or IndexError if it doesn't exist."
ifx<0: x+=self.size
ifx<0: raiseIndexError
forainself.a:
ifx<len(a): returna[x]
x-=len(a)
raiseIndexError

defindex(self, x: T) ->int:
"Count the number of elements < x."
ans=0
forainself.a:
ifa[-1] >=x:
returnans+bisect_left(a, x)
ans+=len(a)
returnans

defindex_right(self, x: T) ->int:
"Count the number of elements <= x."
ans=0
forainself.a:
ifa[-1] >x:
returnans+bisect_right(a, x)
ans+=len(a)
returnans
'''
https://github.com/tatyam-prime/SortedSet
SortedMultiset
add(x) xを追加
discard(x) xを削除(一つだけ)
lt(x) x未満 le(x) x以下 gt(x) xより大きい ge(x) x以上
index(x) x未満の数
index_right(x) x以下の数
s.count(x) xの個数
'''