Get the unheap code in place to fix up the hacky sorting
@@ -213,6 +213,57 @@ def heap_push(heap, row, weight, index, flag):
return 1
+@numba.njit()
+def deheap_sort(heap):
+ """Given an array of heaps (of indices and weights), unpack the heap
+ out to give and array of sorted lists of indices and weights by increasing
+ weight. This is effectively just the second half of heap sort (the first
+ half not being required since we already have the data in a heap).
+
+ Parameters
+ ----------
+ heap : array of shape (3, n_samples, n_neighbors)
+ The heap to turn into sorted lists.
+
+ Returns
+ -------
+ indices, weights: arrays of shape (n_samples, n_neighbors)
+ The indices and weights sorted by increasing weight.
+ """
+ indices = heap[0]
+ weights = heap[1]
+
+ for i in range(indices.shape[0]):
+ heap_end = indices.shape[1] - 1
+ while heap_end >= 0:
+ indices[i, 0], indices[i, heap_end] = \
+ indices[i, heap_end], indices[i, 0]
+ weights[i, 0], weights[i, heap_end] = \
+ weights[i, heap_end], weights[i, 0]
+ heap_end -= 1
+
+ root = 0
+ while root * 2 + 1 < heap_end:
+ left_child = root * 2 + 1
+ right_child = left_child + 1
+ swap = root
+
+ if weights[i, swap] < weights[i, left_child]:
+ swap = left_child
+ if right_child < heap_end and weights[i, swap] < weights[i, right_child]:
+ swap = right_child
+
+ if swap == root:
+ break
+ else:
+ weights[i, root], weights[i, swap] = \
+ weights[i, swap], weights[i, root]
+ indices[i, root], indices[i, swap] = \
+ indices[i, swap], indices[i, root]
+
+ root = swap
+ return indices.astype(np.int64), weights
+
@numba.njit(parallel=True)
def build_candidates(current_graph, n_vertices, n_neighbors, max_candidates,
