Использование qsort в Cython для получения индекса/перестановки сортировки

Обзор

Есть несколько вопросов, похожих на этот, но все они немного отличаются. Чтобы было ясно, если values представляет собой массив целых чисел, я хочу найти perm такое, что sorted_values (values отсортировано некоторым оператором сравнения) задается как

sorted_values[i] = values[perm[i]]

Шаг 1

Как это сделать на Си? Что ж, qsort требует объявления функции сравнения, чтобы сказать вам, больше ли одно значение, чем другое. Если мы сделаем values глобальной переменной, то сможем использовать эту функцию сравнения для сортировки массива perm, изначально установленного в 0:N-1 (где N — длина values), не сравнивая perm[i] с perm[j], а вместо этого сравнивая values[perm[i]] с values[perm[j]]. См. эту ссылку. Некоторый пример кода C:

// sort_test.c
#include <stdio.h>
#include <stdlib.h>

int *values;

int cmpfunc (const void * a, const void * b) {
   return ( values[*(int*)a] - values[*(int*)b] );
}

int main () {
   int i;
   int n = 5;
   int *perm;

   // Assign memory
   values = (int *) malloc(n*sizeof(int));
   perm = (int *) malloc(n*sizeof(int));

   // Set values to random values between 0 and 99
   for (i=0; i<n; i++) values[i] = rand() % 100;

   // Set perm initially to 0:n-1
   for (i=0; i<n; i++) perm[i] = i;

   printf("Before sorting the list is: \n");
   for (i=0; i<n; i++) printf("%d ", values[i]);

   qsort(perm, n, sizeof(int), cmpfunc);

   printf("\nThe sorting permutation is: \n");
   for (i=0; i<n; i++) printf("%d ", perm[i]);

   free(values);
   free(perm);

   printf("\n");

   return(0);
}

Конечно, хитрость заключается в том, чтобы определить values глобально, чтобы cmpfunc мог его видеть.

Шаг 2

Как это сделать на Цитоне? К сожалению, я не могу заставить Cython использовать тот же трюк со значениями, объявленными глобально. Моя лучшая попытка - следующая, основанная на ответе здесь, однако разница в том, что они просто сортируют массив, который им не нужен для индексации/перестановки.

# sort_test_c.pyx
cimport cython
from libc.stdlib cimport qsort

# Try declaring global variable for the sort function
cpdef long[:] values

cdef int cmpfunc (const void *a , const void *b) nogil:
    cdef long a_v = (<long *> a)[0] 
    cdef long b_v = (<long *> b)[0]
    return (values[a_v] - values[b_v]);

def sort(long[:] py_values, long[:] perm, int N):
    # Assign to global
    values = py_values

    # Make sure perm is 0:N-1
    for i in range(N):
        perm[i] = i

    # Perform sort
    qsort(&perm[0], N, perm.strides[0], &cmpfunc)

Это можно скомпилировать с помощью

cythonize -i sort_test_c.pyx

и протестировано со скриптом

# sort_test.py
from sort_test_c import sort
import numpy as np

n = 5
values = np.random.randint(0, 100, n).astype(int)
perm = np.empty(n).astype(int)

sort(values, perm, n)

Однако это жалуется на нашу глобальную переменную values, т.е.

UnboundLocalError: local variable 'values' referenced before assignment
Exception ignored in: 'sort_test_c.cmpfunc

и перестановка сортировки неверна (если только values уже не упорядочены, в этом случае это удача, поскольку perm всегда возвращает массив 0:4). Как я могу это исправить?