более быстрая альтернатива numpy.where?

Question

У меня есть трехмерный массив, заполненный целыми числами от 0 до N. Мне нужен список индексов, соответствующих тому, где массив равен 1, 2, 3, ... N. Я могу сделать это с помощью np.where следующим образом :

N = 300 
shape = (1000,1000,10) 
data = np.random.randint(0,N+1,shape) 
indx = [np.where(data == i_id) for i_id in range(1,data.max()+1)] 

, но это довольно медленно. В соответствии с этим вопросом fast python numpy where functionality? должно быть возможно ускорить поиск индекса довольно много, но я не смог перенести предложенные там методы на проблему получения фактических индексов. Какой был бы лучший способ ускорить выполнение вышеуказанного кода?

В качестве дополнения: Я хочу, чтобы хранить индексы позже, для которых имеет смысл использовать np.ravel_multi_index, чтобы уменьшить размер от экономии 3 индексов только 1, то есть с помощью:

indx = [np.ravel_multi_index(np.where(data == i_id), data.shape) for i_id in range(1, data.max()+1)] 

который ближе к Функция поиска Matlab. Может ли это быть напрямую включено в решение, которое не использует np.where?

Answer 1

Я думаю, что стандартный Векторизованный подход к решению этой проблемы будет в конечный итоге очень большой объем памяти - для Int64 данных, потребуется O (8 * N * data.size) байтов или ~ 22 гигабайта памяти для примера, который вы указали выше. Я предполагаю, что это не вариант.

Вы можете добиться определенного прогресса, используя разреженную матрицу для хранения местоположений уникальных значений. Например:

import numpy as np 
from scipy.sparse import csr_matrix 

def compute_M(data): 
    cols = np.arange(data.size) 
    return csr_matrix((cols, (data.ravel(), cols)), 
         shape=(data.max() + 1, data.size)) 

def get_indices_sparse(data): 
    M = compute_M(data) 
    return [np.unravel_index(row.data, data.shape) for row in M] 

Это имеет преимущество быстрого кода в разреженных матрицах конструктора для организации данных в полезном способе, построением разреженной матрицы, где строка i содержит только индексы, где уплощенные данные равны i.

Чтобы проверить это, я также определить функцию, которая делает ваш простой метод:

def get_indices_simple(data): 
    return [np.where(data == i) for i in range(0, data.max() + 1)] 

Обе функции дают одинаковые результаты для того же входа:

data_small = np.random.randint(0, 100, size=(100, 100, 10)) 
all(np.allclose(i1, i2) 
    for i1, i2 in zip(get_indices_simple(data_small), 
         get_indices_sparse(data_small))) 
# True 

а также разреженный метод на порядок быстрее, чем простой метод для вашего набора данных:

data = np.random.randint(0, 301, size=(1000, 1000, 10)) 

%time ind = get_indices_simple(data) 
# CPU times: user 14.1 s, sys: 638 ms, total: 14.7 s 
# Wall time: 14.8 s 

%time ind = get_indices_sparse(data) 
# CPU times: user 881 ms, sys: 301 ms, total: 1.18 s 
# Wall time: 1.18 s 

%time M = compute_M(data) 
# CPU times: user 216 ms, sys: 148 ms, total: 365 ms 
# Wall time: 363 ms 

Другим преимуществом разреженного метода является то, что матрица M становится очень компактным и эффективным способом хранения всей необходимой информации для последующего использования, как указано в дополнительной части вашего вопроса. Надеюсь, это полезно!

---

Edit: я понял, что ошибка в первоначальной версии: он не если какое-либо значение в диапазоне не появляется в данном: что теперь закрепился выше.

Answer 2

Вот один Векторизованный подход -

# Mask of matches for data elements against all IDs from 1 to data.max() 
mask = data == np.arange(1,data.max()+1)[:,None,None,None] 

# Indices of matches across all IDs and their linear indices 
idx = np.argwhere(mask.reshape(N,-1)) 

# Get cut indices where IDs shift 
_,cut_idx = np.unique(idx[:,0],return_index=True) 

# Cut at shifts to give us the final indx output 
out = np.hsplit(idx[:,1],cut_idx[1:]) 

Answer 3

Я обдумывал это и понял, что существует более интуитивный (но немного более медленный) подход к решению этого вопроса с использованием Pandas groupby().Рассмотрю это:

import numpy as np 
import pandas as pd 

def get_indices_pandas(data): 
    d = data.ravel() 
    f = lambda x: np.unravel_index(x.index, data.shape) 
    return pd.Series(d).groupby(d).apply(f) 

Это возвращает тот же результат, как get_indices_simple из моего предыдущего ответа:

data_small = np.random.randint(0, 100, size=(100, 100, 10)) 
all(np.allclose(i1, i2) 
    for i1, i2 in zip(get_indices_simple(data_small), 
         get_indices_pandas(data_small))) 
# True 

И это Панда подход лишь немного медленнее, чем меньше интуитивный подход матрицы:

data = np.random.randint(0, 301, size=(1000, 1000, 10)) 

%time ind = get_indices_simple(data) 
# CPU times: user 14.2 s, sys: 665 ms, total: 14.8 s 
# Wall time: 14.9 s 

%time ind = get_indices_sparse(data) 
# CPU times: user 842 ms, sys: 277 ms, total: 1.12 s 
# Wall time: 1.12 s 

%time ind = get_indices_pandas(data) 
# CPU times: user 1.16 s, sys: 326 ms, total: 1.49 s 
# Wall time: 1.49 s 

Answer 4

В основном, большинство ответов на другой вопрос имеют сообщение «использовать косвенную сортировку».

Мы можем получить линейные индексы (так похожие на find в MATLAB), соответствующая i = [0..N] с вызовом numpy.argsort над сплющенным массивом:

flat = data.ravel() 
lin_idx = np.argsort(flat, kind='mergesort') 

Но тогда мы получаем один большой массив; какие индексы принадлежат i? Мы просто разделить массив индексов на основе подсчетов каждого i:

ans = np.split(lin_idx, np.cumsum(np.bincount(flat)[:-1])) 

Если вам еще нужны 3D-индексы где-то, вы можете использовать numpy.unravel_index.