Python Indexed Priority Queue O(logN) seat(), O(logN) leave()
Inspired by William Fiset\'s binary heap indexed priority queue implementation:\n\nhttps://www.youtube.com/watch?v=jND_WJ8r7FE&t=342s\n\nAlso took some notes fr
Inspired by William Fiset\'s binary heap indexed priority queue implementation:\n\nhttps://www.youtube.com/watch?v=jND_WJ8r7FE&t=342s\n\nAlso took some notes from @spatar \'s implementation:\n\nhttps://leetcode.com/problems/exam-room/discuss/139941/Python-O(log-n)-time-for-both-seat()-and-leave()-with-heapq-and-dicts-Detailed-explanation\n\nHowever, as this is a true indexed priority queue the heap only stores valid segments. Invalid segments are eagerly removed in O(logn) as well by swapping with the value at the bottom right of the heap, removing it O(1), and updating the heap invariant O(logn). One downside is for really large inputs, the index priority queue must construct a key map first. I might try and rewrite the priority queue so this isn\'t required, but i\'m not sure if that\'s possible. \n\nAnyways, upvote if you like. Despite getting TLE on LC this is the most comprehensive way to solve this problem.\n\n```\nclass IPQ:\n    \n    def __init__(self, N):\n        self.values = [None] * N\n        # each index of the position map is the key index\n        # of an object. the value at that index is the \n        # top-left->bottom-right index of the object in the\n        # heap\n        # position map is key_index:heap_pos\n        self.pm = [-1] * N\n        # inverse map is heap_pos:key_index\n        self.im = [-1] * N\n        self.sz = 0\n        \n    def get_heap(self):\n        """useful for debugging and understanding"""\n        return [self.values[i] for i in self.im if i != -1]\n        \n    def _swim(self, h_i):\n        """\n        Time O(logn)\n        """\n        p = (h_i-1)//2\n        # if the idx of the current node is not the root node\n        # (which has idx=0) and the value of the current node\n        # is less than the value of the parent node...\n        while h_i > 0 and self._less(h_i, p):\n            # we need to swap the current node with the parent node\n            self._swap(h_i, p)\n            h_i = p\n            p = (h_i-1)//2\n\n    def _sink(self, h_i):\n        """\n        Time O(logn)\n        Sinks the node at index i by swapping\n        itself with the smallest of the left\n        or the right child node.\n        """\n        while True:\n            left = 2*h_i+1\n            right = 2*h_i+2\n            smallest = left\n\n            if right < self.sz and self._less(right, left):\n                smallest = right\n\n            if left >= self.sz or self._less(h_i, smallest):\n                break\n\n            self._swap(smallest, h_i)\n            h_i = smallest\n\n    def _swap(self, h_i, h_j):\n        """Time O(1)"""\n        # use inverse map to update position map\n        # heap pos j-> ki j -> pos j = i\n        self.pm[self.im[h_j]] = h_i\n        self.pm[self.im[h_i]] = h_j\n        # update inverse map\n        self.im[h_i], self.im[h_j] = self.im[h_j], self.im[h_i]\n\n    def _less(self, h_i, h_j):\n        """\n        Time O(1)\n        Takes heap positions i and j, looks up values and returns\n        if heap element i < heap element j\n        """\n        return self.values[self.im[h_i]] < self.values[self.im[h_j]]\n\n    def insert(self, ki, value):\n        """Time O(logn)"""\n        \n        # update values map\n        self.values[ki] = value\n        # update heap position map\n        self.pm[ki] = self.sz\n        # update inverse position (heap_pos:ki) map\n        self.im[self.sz] = ki\n        # because we start with inserted val at end, we\n        # swim it up\n        self._swim(self.sz)\n        # increment size\n        self.sz += 1\n                \n    def remove(self, ki):\n        """Time O(logn)\n        Remove a segment that starts with ki\n        """\n        # first find the heap index of that key\n        h_i = self.pm[ki]\n        # decrement the size\n        self.sz -= 1\n        # next swap that key index with the last thing in the heap\n        self._swap(h_i, self.sz)\n        # because i was swapped, i now is the position\n        # of the node that was previously at the bottom\n        # right of the heap.  we don\'t know if we need to\n        # push that node down or bring it closer to the top\n        # so we call sink(i) and swim(i) to make sure it is \n        # in the right place.\n        self._sink(h_i)\n        self._swim(h_i)\n\n        val = self.values[ki]\n        self.values[ki] = None\n        self.pm[ki] = -1\n        self.im[self.sz] = -1\n        return ki, val\n\n    def update(self, ki, value):\n        """Time O(logn)"""\n        h_i = self.pm[ki]\n        self.values[ki] = value\n        self._sink(h_i)\n        self._swim(h_i)\n        \n    def poll(self):\n        """\n        Time O(logn + k)\n        where k is num of vals at top of heap\n        with same min val\n        \n        guaranteed to have at least 1 item in heap\n        """\n        # find ki (start idx) with greatest val\n        # start from h_i = 0 and go down until val increases\n        # top of heap\n        ki = self.im[0]\n        val = self.values[ki]\n        queue = deque([0])\n        while queue:\n            h_i = queue.popleft()\n            left = 2*h_i+1\n            right = 2*h_i+2\n            if right < self.sz and self.values[self.im[right]] == val:\n                queue.append(right)\n            if left < self.sz and self.values[self.im[left]] == val:\n                queue.append(left)\n            ki = min(ki, self.im[h_i])\n        return self.remove(ki)\n\nclass ExamRoom:\n\n    def __init__(self, n: int):\n\n        self.n = n\n        # we use this to check a cornercase where \n        # the seat leaving opens up a segment of \n        # start and end = that seat index\n        self.seats = [0]*n\n        \n        # initialize queue\n        self.ipq = IPQ(n)\n        \n        # these are just lookup keys for the priority queue\n        # we just need n number of unique keys, their values are\n        # irrelevant\n        self.availKeys = set(range(n))\n        \n        # we need a way to remove a priority queue lookup key\n        # based on just the upper or lower segment bound\n        \n        # lower_idx:key\n        self.leftKey = {}\n        # upper_idx:key\n        self.rightKey = {}\n        \n        # push first segment to queue\n        self.put_segment(0,n-1)\n        \n    def clean_key_map(self, key, seg):\n        if key == None:\n            return\n\n        self.availKeys.add(key)\n        del self.leftKey[seg[1]]\n        del self.rightKey[seg[2]]\n        \n    def put_segment(self, first, last):\n        if first == 0 or last == self.n - 1:\n            priority = last-first\n        else:\n            priority = (last - first) // 2\n            \n        segment = [-priority, first, last]\n        \n        key = self.availKeys.pop()\n        self.leftKey[first] = key\n        self.rightKey[last] = key\n        self.ipq.insert(key, segment)\n\n    def seat(self) -> int:\n        k, (p, first, last) = self.ipq.poll()\n        self.clean_key_map(k, (p,first,last))\n\n        if first == 0:\n            out = 0\n            if first != last:\n                self.put_segment(first+1, last)\n        \n        elif last == self.n-1:\n            out = last\n            if first != last:\n                self.put_segment(first, last-1)\n        else:\n            out = first + (last - first) // 2\n            \n            if out > first:\n                self.put_segment(first, out-1)\n            if out < last:\n                self.put_segment(out+1, last)\n\n        self.seats[out] = 1\n\n        return out\n\n    def leave(self, p: int) -> None:\n        lowerAvail, upperAvail = True, True\n        \n        # check to see if there is\n        # avail seg that ends on left and right sides\n        if p == 0 or self.seats[p-1]:\n            lowerAvail = False\n        if p+1==self.n or self.seats[p+1]:\n            upperAvail = False\n        \n        self.seats[p] = 0 \n\n        # there is no space on either side\n        # make avail segment of size 1\n        if not lowerAvail and not upperAvail:\n            self.put_segment(p,p)\n            return\n\n        if lowerAvail:\n            # get priority queue key for segment which ends at left upper\n            leftKey = self.rightKey[p-1]\n            _, leftSeg = self.ipq.remove(leftKey)\n            self.clean_key_map(leftKey, leftSeg)\n        if upperAvail:\n            # get priority queue key for segment which starts at right lower\n            rightKey = self.leftKey[p+1]\n            _, rightSeg = self.ipq.remove(rightKey)\n            self.clean_key_map(rightKey, rightSeg)\n\n        # we have both lower and upper segments\n        if lowerAvail and upperAvail:\n            self.put_segment(leftSeg[1], rightSeg[2])\n        # there is only a lower segment\n        elif lowerAvail:\n            self.put_segment(leftSeg[1], p)\n        # there is only an upper segment\n        elif upperAvail:\n            self.put_segment(p, rightSeg[2])
['Heap (Priority Queue)']