1"""
2Originally part of CellProfiler, code licensed under both GPL and BSD
3licenses.
4Website: http://www.cellprofiler.org
5
6Copyright (c) 2003-2009 Massachusetts Institute of Technology
7Copyright (c) 2009-2011 Broad Institute
8All rights reserved.
9
10Original author: Lee Kamentsky
11"""
12
13from libc.stdlib cimport free, malloc, realloc
14
15
16cdef struct Heap:
17    Py_ssize_t items
18    Py_ssize_t space
19    Heapitem *data
20    Heapitem **ptrs
21
22cdef inline Heap *heap_from_numpy2() nogil:
23    cdef Py_ssize_t k
24    cdef Heap *heap
25    heap = <Heap *> malloc(sizeof (Heap))
26    heap.items = 0
27    heap.space = 1000
28    heap.data = <Heapitem *> malloc(heap.space * sizeof(Heapitem))
29    heap.ptrs = <Heapitem **> malloc(heap.space * sizeof(Heapitem *))
30    for k in range(heap.space):
31        heap.ptrs[k] = heap.data + k
32    return heap
33
34cdef inline void heap_done(Heap *heap) nogil:
35   free(heap.data)
36   free(heap.ptrs)
37   free(heap)
38
39cdef inline void swap(Py_ssize_t a, Py_ssize_t b, Heap *h) nogil:
40    h.ptrs[a], h.ptrs[b] = h.ptrs[b], h.ptrs[a]
41
42
43######################################################
44# heappop - inlined
45#
46# pop an element off the heap, maintaining heap invariant
47#
48# Note: heap ordering is the same as python heapq, i.e., smallest first.
49######################################################
50cdef inline void heappop(Heap *heap, Heapitem *dest) nogil:
51
52    cdef Py_ssize_t i, smallest, l, r # heap indices
53
54    #
55    # Start by copying the first element to the destination
56    #
57    dest[0] = heap.ptrs[0][0]
58    heap.items -= 1
59
60    # if the heap is now empty, we can return, no need to fix heap.
61    if heap.items == 0:
62        return
63
64    #
65    # Move the last element in the heap to the first.
66    #
67    swap(0, heap.items, heap)
68
69    #
70    # Restore the heap invariant.
71    #
72    i = 0
73    smallest = i
74    while True:
75        # loop invariant here: smallest == i
76
77        # find smallest of (i, l, r), and swap it to i's position if necessary
78        l = i * 2 + 1 #__left(i)
79        r = i * 2 + 2 #__right(i)
80        if l < heap.items:
81            if smaller(heap.ptrs[l], heap.ptrs[i]):
82                smallest = l
83            if r < heap.items and smaller(heap.ptrs[r], heap.ptrs[smallest]):
84                smallest = r
85        else:
86            # this is unnecessary, but trims 0.04 out of 0.85 seconds...
87            break
88        # the element at i is smaller than either of its children, heap
89        # invariant restored.
90        if smallest == i:
91                break
92        # swap
93        swap(i, smallest, heap)
94        i = smallest
95
96##################################################
97# heappush - inlined
98#
99# push the element onto the heap, maintaining the heap invariant
100#
101# Note: heap ordering is the same as python heapq, i.e., smallest first.
102##################################################
103cdef inline void heappush(Heap *heap, Heapitem *new_elem) nogil:
104
105    cdef Py_ssize_t child = heap.items
106    cdef Py_ssize_t parent
107    cdef Py_ssize_t k
108    cdef Heapitem *new_data
109
110    # grow if necessary
111    if heap.items == heap.space:
112      heap.space = heap.space * 2
113      new_data = <Heapitem*>realloc(<void*>heap.data,
114                    <Py_ssize_t>(heap.space * sizeof(Heapitem)))
115      heap.ptrs = <Heapitem**>realloc(<void*>heap.ptrs,
116                    <Py_ssize_t>(heap.space * sizeof(Heapitem *)))
117      for k in range(heap.items):
118          heap.ptrs[k] = new_data + (heap.ptrs[k] - heap.data)
119      for k in range(heap.items, heap.space):
120          heap.ptrs[k] = new_data + k
121      heap.data = new_data
122
123    # insert new data at child
124    heap.ptrs[child][0] = new_elem[0]
125    heap.items += 1
126
127    # restore heap invariant, all parents <= children
128    while child > 0:
129        parent = (child + 1) // 2 - 1 # __parent(i)
130
131        if smaller(heap.ptrs[child], heap.ptrs[parent]):
132            swap(parent, child, heap)
133            child = parent
134        else:
135            break
136