Aluode/PerceptionLabPortable

1import collections
2
3import numpy as np
4
5from numba.core import types, config
6
7
8QuicksortImplementation = collections.namedtuple(
9    'QuicksortImplementation',
10    (# The compile function itself
11     'compile',
12     # All subroutines exercised by test_sort
13     'partition', 'partition3', 'insertion_sort',
14     # The top-level function
15     'run_quicksort',
16     ))
17
18
19Partition = collections.namedtuple('Partition', ('start', 'stop'))
20
21# Under this size, switch to a simple insertion sort
22SMALL_QUICKSORT = 15
23
24MAX_STACK = 100
25
26
27def make_quicksort_impl(wrap, lt=None, is_argsort=False, is_list=False, is_np_array=False):
28
29    if config.USE_LEGACY_TYPE_SYSTEM:
30        intp = types.intp
31    else:
32        intp = types.py_int
33    zero = intp(0)
34
35    # Two subroutines to make the core algorithm generic wrt. argsort
36    # or normal sorting.  Note the genericity may make basic sort()
37    # slightly slower (~5%)
38    if is_argsort:
39        if is_list:
40            @wrap
41            def make_res(A):
42                return [x for x in range(len(A))]
43        else:
44            @wrap
45            def make_res(A):
46                return np.arange(A.size)
47
48        @wrap
49        def GET(A, idx_or_val):
50            return A[idx_or_val]
51
52    else:
53        @wrap
54        def make_res(A):
55            return A
56
57        @wrap
58        def GET(A, idx_or_val):
59            return idx_or_val
60
61    def default_lt(a, b):
62        """
63        Trivial comparison function between two keys.
64        """
65        return a < b
66
67    LT = wrap(lt if lt is not None else default_lt)
68
69    @wrap
70    def insertion_sort(A, R, low, high):
71        """
72        Insertion sort A[low:high + 1]. Note the inclusive bounds.
73        """
74        assert low >= 0
75        if high <= low:
76            return
77
78        for i in range(low + 1, high + 1):
79            k = R[i]
80            v = GET(A, k)
81            # Insert v into A[low:i]
82            j = i
83            while j > low and LT(v, GET(A, R[j - 1])):
84                # Make place for moving A[i] downwards
85                R[j] = R[j - 1]
86                j -= 1
87            R[j] = k
88
89    @wrap
90    def partition(A, R, low, high):
91        """
92        Partition A[low:high + 1] around a chosen pivot.  The pivot's index
93        is returned.
94        """
95        assert low >= 0
96        assert high > low
97
98        mid = (low + high) >> 1
99        # NOTE: the pattern of swaps below for the pivot choice and the
100        # partitioning gives good results (i.e. regular O(n log n))
101        # on sorted, reverse-sorted, and uniform arrays.  Subtle changes
102        # risk breaking this property.
103
104        # median of three {low, middle, high}
105        if LT(GET(A, R[mid]), GET(A, R[low])):
106            R[low], R[mid] = R[mid], R[low]
107        if LT(GET(A, R[high]), GET(A, R[mid])):
108            R[high], R[mid] = R[mid], R[high]
109        if LT(GET(A, R[mid]), GET(A, R[low])):
110            R[low], R[mid] = R[mid], R[low]
111        pivot = GET(A, R[mid])
112
113        # Temporarily stash the pivot at the end
114        R[high], R[mid] = R[mid], R[high]
115        i = low
116        j = high - 1
117        while True:
118            while i < high and LT(GET(A, R[i]), pivot):
119                i += 1
120            while j >= low and LT(pivot, GET(A, R[j])):
121                j -= 1
122            if i >= j:
123                break
124            R[i], R[j] = R[j], R[i]
125            i += 1
126            j -= 1
127        # Put the pivot back in its final place (all items before `i`
128        # are smaller than the pivot, all items at/after `i` are larger)
129        R[i], R[high] = R[high], R[i]
130        return i
131
132    @wrap
133    def partition3(A, low, high):
134        """
135        Three-way partition [low, high) around a chosen pivot.
136        A tuple (lt, gt) is returned such that:
137            - all elements in [low, lt) are < pivot
138            - all elements in [lt, gt] are == pivot
139            - all elements in (gt, high] are > pivot
140        """
141        mid = (low + high) >> 1
142        # median of three {low, middle, high}
143        if LT(A[mid], A[low]):
144            A[low], A[mid] = A[mid], A[low]
145        if LT(A[high], A[mid]):
146            A[high], A[mid] = A[mid], A[high]
147        if LT(A[mid], A[low]):
148            A[low], A[mid] = A[mid], A[low]
149        pivot = A[mid]
150
151        A[low], A[mid] = A[mid], A[low]
152        lt = low
153        gt = high
154        i = low + 1
155        while i <= gt:
156            if LT(A[i], pivot):
157                A[lt], A[i] = A[i], A[lt]
158                lt += 1
159                i += 1
160            elif LT(pivot, A[i]):
161                A[gt], A[i] = A[i], A[gt]
162                gt -= 1
163            else:
164                i += 1
165        return lt, gt
166
167    @wrap
168    def run_quicksort1(A):
169        R = make_res(A)
170
171        if len(A) < 2:
172            return R
173
174        stack = [Partition(zero, zero)] * MAX_STACK
175        stack[0] = Partition(zero, len(A) - 1)
176        n = 1
177
178        while n > 0:
179            n -= 1
180            low, high = stack[n]
181            # Partition until it becomes more efficient to do an insertion sort
182            while high - low >= SMALL_QUICKSORT:
183                assert n < MAX_STACK
184                i = partition(A, R, low, high)
185                # Push largest partition on the stack
186                if high - i > i - low:
187                    # Right is larger
188                    if high > i:
189                        stack[n] = Partition(i + 1, high)
190                        n += 1
191                    high = i - 1
192                else:
193                    if i > low:
194                        stack[n] = Partition(low, i - 1)
195                        n += 1
196                    low = i + 1
197
198            insertion_sort(A, R, low, high)
199
200        return R
201
202    if is_np_array:
203        @wrap
204        def run_quicksort(A):
205            if A.ndim == 1:
206                return run_quicksort1(A)
207            else:
208                for idx in np.ndindex(A.shape[:-1]):
209                    run_quicksort1(A[idx])
210                return A
211    else:
212        @wrap
213        def run_quicksort(A):
214            return run_quicksort1(A)
215
216    # Unused quicksort implementation based on 3-way partitioning; the
217    # partitioning scheme turns out exhibiting bad behaviour on sorted arrays.
218    @wrap
219    def _run_quicksort(A):
220        stack = [Partition(zero, zero)] * 100
221        stack[0] = Partition(zero, len(A) - 1)
222        n = 1
223
224        while n > 0:
225            n -= 1
226            low, high = stack[n]
227            # Partition until it becomes more efficient to do an insertion sort
228            while high - low >= SMALL_QUICKSORT:
229                assert n < MAX_STACK
230                l, r = partition3(A, low, high)
231                # One trivial (empty) partition => iterate on the other
232                if r == high:
233                    high = l - 1
234                elif l == low:
235                    low = r + 1
236                # Push largest partition on the stack
237                elif high - r > l - low:
238                    # Right is larger
239                    stack[n] = Partition(r + 1, high)
240                    n += 1
241                    high = l - 1
242                else:
243                    stack[n] = Partition(low, l - 1)
244                    n += 1
245                    low = r + 1
246
247            insertion_sort(A, low, high)
248
249
250    return QuicksortImplementation(wrap,
251                                   partition, partition3, insertion_sort,
252                                   run_quicksort)
253
254
255def make_py_quicksort(*args, **kwargs):
256    return make_quicksort_impl((lambda f: f), *args, **kwargs)
257
258def make_jit_quicksort(*args, **kwargs):
259    from numba.core.extending import register_jitable
260    return make_quicksort_impl((lambda f: register_jitable(f)),
261                               *args, **kwargs)
262