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
|
# cython: boundscheck=False, wraparound=False, cdivision=True
import numpy as np
ctypedef struct Pair:
int key
double value
cdef class Heap:
"""Binary heap.
Heap stores values and keys. Values are passed explicitly, whereas keys
are assigned implicitly to natural numbers (from 0 to n - 1).
The supported operations (all have O(log n) time complexity):
* Return the current minimum value and the corresponding key.
* Remove the current minimum value.
* Change the value of the given key. Note that the key must be still
in the heap.
The heap is stored as an array, where children of parent i have indices
2 * i + 1 and 2 * i + 2. All public methods are based on `sift_down` and
`sift_up` methods, which restore the heap property by moving an element
down or up in the heap.
"""
cdef int[:] index_by_key
cdef int[:] key_by_index
cdef double[:] values
cdef int size
def __init__(self, double[:] values):
self.size = values.shape[0]
self.index_by_key = np.arange(self.size, dtype=np.intc)
self.key_by_index = np.arange(self.size, dtype=np.intc)
self.values = values.copy()
cdef int i
# Create the heap in a linear time. The algorithm sequentially sifts
# down items starting from lower levels.
for i in reversed(range(self.size / 2)):
self.sift_down(i)
cpdef Pair get_min(self):
return Pair(self.key_by_index[0], self.values[0])
cpdef void remove_min(self):
self.swap(0, self.size - 1)
self.size -= 1
self.sift_down(0)
cpdef void change_value(self, int key, double value):
cdef int index = self.index_by_key[key]
cdef double old_value = self.values[index]
self.values[index] = value
if value < old_value:
self.sift_up(index)
else:
self.sift_down(index)
cdef void sift_up(self, int index):
cdef int parent = Heap.parent(index)
while index > 0 and self.values[parent] > self.values[index]:
self.swap(index, parent)
index = parent
parent = Heap.parent(index)
cdef void sift_down(self, int index):
cdef int child = Heap.left_child(index)
while child < self.size:
if (child + 1 < self.size and
self.values[child + 1] < self.values[child]):
child += 1
if self.values[index] > self.values[child]:
self.swap(index, child)
index = child
child = Heap.left_child(index)
else:
break
@staticmethod
cdef inline int left_child(int parent):
return (parent << 1) + 1
@staticmethod
cdef inline int parent(int child):
return (child - 1) >> 1
cdef void swap(self, int i, int j):
self.values[i], self.values[j] = self.values[j], self.values[i]
cdef int key_i = self.key_by_index[i]
cdef int key_j = self.key_by_index[j]
self.key_by_index[i] = key_j
self.key_by_index[j] = key_i
self.index_by_key[key_i] = j
self.index_by_key[key_j] = i
|
