fromsortedcontainersimportSortedListclassExamRoom:def__init__(self,n:int):defdist(x):l,r=xreturnr-l-1ifl==-1orr==nelse(r-l)>>1self.n=nself.ts=SortedList(key=lambdax:(-dist(x),x[0]))self.left={}self.right={}self.add((-1,n))defseat(self)->int:s=self.ts[0]p=(s[0]+s[1])>>1ifs[0]==-1:p=0elifs[1]==self.n:p=self.n-1self.delete(s)self.add((s[0],p))self.add((p,s[1]))returnpdefleave(self,p:int)->None:l,r=self.left[p],self.right[p]self.delete((l,p))self.delete((p,r))self.add((l,r))defadd(self,s):self.ts.add(s)self.left[s[1]]=s[0]self.right[s[0]]=s[1]defdelete(self,s):self.ts.remove(s)self.left.pop(s[1])self.right.pop(s[0])# Your ExamRoom object will be instantiated and called as such:# obj = ExamRoom(n)# param_1 = obj.seat()# obj.leave(p)
classExamRoom{privateTreeSet<int[]>ts=newTreeSet<>((a,b)->{intd1=dist(a),d2=dist(b);returnd1==d2?a[0]-b[0]:d2-d1;});privateMap<Integer,Integer>left=newHashMap<>();privateMap<Integer,Integer>right=newHashMap<>();privateintn;publicExamRoom(intn){this.n=n;add(newint[]{-1,n});}publicintseat(){int[]s=ts.first();intp=(s[0]+s[1])>>1;if(s[0]==-1){p=0;}elseif(s[1]==n){p=n-1;}del(s);add(newint[]{s[0],p});add(newint[]{p,s[1]});returnp;}publicvoidleave(intp){intl=left.get(p),r=right.get(p);del(newint[]{l,p});del(newint[]{p,r});add(newint[]{l,r});}privateintdist(int[]s){intl=s[0],r=s[1];returnl==-1||r==n?r-l-1:(r-l)>>1;}privatevoidadd(int[]s){ts.add(s);left.put(s[1],s[0]);right.put(s[0],s[1]);}privatevoiddel(int[]s){ts.remove(s);left.remove(s[1]);right.remove(s[0]);}}/** * Your ExamRoom object will be instantiated and called as such: * ExamRoom obj = new ExamRoom(n); * int param_1 = obj.seat(); * obj.leave(p); */
intN;intdist(constpair<int,int>&p){auto[l,r]=p;if(l==-1||r==N)returnr-l-1;return(r-l)>>1;}structcmp{booloperator()(constpair<int,int>&a,constpair<int,int>&b)const{intd1=dist(a),d2=dist(b);returnd1==d2?a.first<b.first:d1>d2;};};classExamRoom{public:ExamRoom(intn){N=n;this->n=n;add({-1,n});}intseat(){autos=*ts.begin();intp=(s.first+s.second)>>1;if(s.first==-1){p=0;}elseif(s.second==n){p=n-1;}del(s);add({s.first,p});add({p,s.second});returnp;}voidleave(intp){intl=left[p],r=right[p];del({l,p});del({p,r});add({l,r});}private:set<pair<int,int>,cmp>ts;unordered_map<int,int>left;unordered_map<int,int>right;intn;voidadd(pair<int,int>s){ts.insert(s);left[s.second]=s.first;right[s.first]=s.second;}voiddel(pair<int,int>s){ts.erase(s);left.erase(s.second);right.erase(s.first);}};/** * Your ExamRoom object will be instantiated and called as such: * ExamRoom* obj = new ExamRoom(n); * int param_1 = obj->seat(); * obj->leave(p); */
typeExamRoomstruct{rbt*redblacktree.Treeleftmap[int]intrightmap[int]intnint}funcConstructor(nint)ExamRoom{dist:=func(s[]int)int{ifs[0]==-1||s[1]==n{returns[1]-s[0]-1}return(s[1]-s[0])>>1}cmp:=func(a,bany)int{x,y:=a.([]int),b.([]int)d1,d2:=dist(x),dist(y)ifd1==d2{returnx[0]-y[0]}returnd2-d1}this:=ExamRoom{redblacktree.NewWith(cmp),map[int]int{},map[int]int{},n}this.add([]int{-1,n})returnthis}func(this*ExamRoom)Seat()int{s:=this.rbt.Left().Key.([]int)p:=(s[0]+s[1])>>1ifs[0]==-1{p=0}elseifs[1]==this.n{p=this.n-1}this.del(s)this.add([]int{s[0],p})this.add([]int{p,s[1]})returnp}func(this*ExamRoom)Leave(pint){l,_:=this.left[p]r,_:=this.right[p]this.del([]int{l,p})this.del([]int{p,r})this.add([]int{l,r})}func(this*ExamRoom)add(s[]int){this.rbt.Put(s,struct{}{})this.left[s[1]]=s[0]this.right[s[0]]=s[1]}func(this*ExamRoom)del(s[]int){this.rbt.Remove(s)delete(this.left,s[1])delete(this.right,s[0])}/** * Your ExamRoom object will be instantiated and called as such: * obj := Constructor(n); * param_1 := obj.Seat(); * obj.Leave(p); */