Design a Number Container System LeetCode solution in Python

1
class NumberContainers:
2
    def __init__(self):
3
        self.numbersByIndex = {}
4
        self.numberIndexes = defaultdict(set)
5
        self.numberIndexesHeap = defaultdict(list)
6

7
    def change(self, index: int, number: int) -> None:
8
        if index in self.numbersByIndex:
9
            if number != self.numbersByIndex[index]:
10
                self.numberIndexes[self.numbersByIndex[index]].remove(index)
11
                self.numbersByIndex[index] = number
12
                self.numberIndexes[number].add(index)
13
                heappush(self.numberIndexesHeap[number], index)
14
        else:
15
            self.numbersByIndex[index] = number
16
            self.numberIndexes[number].add(index)
17
            heappush(self.numberIndexesHeap[number], index)
18

19
    def find(self, number: int) -> int:
20
        while (
21
            self.numberIndexesHeap[number]
22
            and self.numberIndexesHeap[number][0] not in self.numberIndexes[number]
23
        ):
24
            heappop(
25
                self.numberIndexesHeap[number]
26
            )  # make sure the smallest index in heap is still an index for number
27
        return (
28
            self.numberIndexesHeap[number][0] if self.numberIndexesHeap[number] else -1
29
        )