Rubiks/indexqueue.py at main · Zingler/Rubiks · 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
class IndexedQueue:
    def __init__(self):
        self.heap = [()] # Null element at index 0
        self.index = {}

    def _add(self, priority, object):
        self.heap.append((priority, object))
        index = len(self.heap) - 1
        self.index[object] = index
        self._siftup(index)

    def add_or_update(self, priority, object):
        existing = self.get(object)
        if not existing:
            self._add(priority, object)
        else:
            index = self.index[object]
            self.heap[index] = (priority, object)

            if priority < existing[0]:
                self._siftup(index)
            else:
                self._siftdown(index)


    def get(self, object):
        index = self.index.get(object)
        if index:
            return self.heap[index]
        else:
            return None

    def size(self):
        return len(self.heap) - 1

    def empty(self):
        return self.size() == 0

    def pop(self):
        last_index = len(self.heap) - 1
        self._swap(1, last_index)

        ret_val = self.heap.pop()
        if not self.empty():
            self._siftdown(1)


        del self.index[ret_val[1]]
        return ret_val

    def _swap(self, i, j):
        a = self.heap[i]
        b = self.heap[j]
        self.heap[i], self.heap[j] = b, a
        self.index[a[1]], self.index[b[1]] = j, i


    @staticmethod
    def parent(index):
        return index // 2

    @staticmethod
    def left(index):
        return index * 2

    @staticmethod
    def right(index):
        return index * 2 + 1


    def _siftup(self, index):
        if index == 1:
            return
        o = self.heap[index]
        p_i = IndexedQueue.parent(index)
        p = self.heap[p_i]

        if o[0] < p[0]:
            self._swap(index, p_i)
            self._siftup(p_i)

    def _siftdown(self, index):
        left_i = IndexedQueue.left(index)
        right_i = IndexedQueue.right(index)
        left = None
        right = None
        parent = self.heap[index]

        if left_i < len(self.heap):
            left = self.heap[left_i]
        if right_i < len(self.heap):
            right = self.heap[right_i]

        if left == None and right == None:
            return
        elif right == None:
            if left[0] < parent[0]:
                self._swap(index, left_i)
                self._siftdown(left_i)
        else:
            if left[0] < right[0]:
                if left[0] < parent[0]:
                    self._swap(index, left_i)
                    self._siftdown(left_i)
            else:
                if right[0] < parent[0]:
                    self._swap(index, right_i)
                    self._siftdown(right_i)

if __name__ == "__main__":
    q = IndexedQueue()
    q.add_or_update(9, "c")
    q.add_or_update(2, "b")
    q.add_or_update(5, "d")
    q.add_or_update(1, "a")
    q.add_or_update(3, "c")
    q.add_or_update(0, "!")
    q.add_or_update(9, "!")

    while not q.empty():
        print(q.pop())