Нахождение индексов различных элементов в векторизации

Question

У меня есть список int с, a, от 0 до 3000. len(a) = 3000. У меня есть for l oop, который перебирает этот список в поисках индексов каждого элемента в большем массиве.

import numpy as np

a = [i for i in range(3000)]
array = np.random.randint(0, 3000, size(12, 1000, 1000))
newlist = []

for i in range(0, len(a)):
    coord = np.where(array == list[i])
    newlist.append(coord)

Как видите, coord будет 3 массивами координат x, y, z для значений в трехмерной матрице, которые равны значению в списке.

Есть ли способ сделать это в векторизации без for l oop?

Выходными данными должен быть список кортежей, по одному для каждого элемента в a :

# each coord looks like this:
print(coord)
(array[1, ..., 1000], array[2, ..., 1000], array[2, ..., 12])

# combined over all the iterations:
print(newlist)
[coord1, coord2, ..., coord3000]

Answer 1

Существует фактически полностью векторизованное решение, несмотря на тот факт, что все полученные массивы имеют разные размеры. Идея такова:

1. Сортировка всех элементов массива вместе с их координатами. argsort идеально подходит для такого рода вещей.
2. Найдите точки среза в отсортированных данных, чтобы вы знали, где разбить массив, например, с помощью diff и flatnonzero.
3. split массив координат по найденным вами индексам. Если у вас отсутствуют элементы, вам может потребоваться сгенерировать ключ на основе первого элемента каждого прогона.

Вот пример, который поможет вам пройти через него. Допустим, у вас есть d -мерный массив с размером n. Ваши координаты будут (d, n) массивом:

d = arr.ndim
n = arr.size

. Вы можете сгенерировать координатные массивы напрямую с помощью np.indices:

coords = np.indices(arr.shape)

Now ravel / reshape данные и координаты в массив (n,) и (d, n) соответственно:

arr = arr.ravel()  # Ravel guarantees C-order no matter the source of the data
coords = coords.reshape(d, n)  # C-order by default as a result of `indices` too

Теперь сортируйте данные:

order = np.argsort(arr)
arr = arr[order]
coords = coords[:, order]

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

change = np.diff(arr, prepend=arr[0] - 1)

Индексы местоположений дают точки останова в массиве:

locs = np.flatnonzero(change)

Теперь вы можете разделить данные в следующих местах:

result = np.split(coords, locs[1:], axis=1)

И вы можете создать ключ фактически найденных значений:

key = arr[locs]

Если вы очень уверены, что все значения присутствуют в массиве, тогда вам не нужен ключ. Вместо этого вы можете вычислить locs как просто np.diff(arr) и result как просто np.split(coords, inds, axis=1).

Каждый элемент в result уже соответствует индексации, используемой where / nonzero, но в виде массива numpy. Если вам нужен кортеж, вы можете сопоставить его с кортежем:

result = [tuple(inds) for inds in result]

TL; DR

Объединение всего этого в функцию:

def find_locations(arr):
    coords = np.indices(arr.shape).reshape(arr.ndim, arr.size)
    arr = arr.ravel()
    order = np.argsort(arr)
    arr = arr[order]
    coords = coords[:, order]
    locs = np.flatnonzero(np.diff(arr, prepend=arr[0] - 1))
    return arr[locs], np.split(coords, locs[1:], axis=1)

Вы можете вернуть список индексных массивов с пустыми массивами для отсутствующих элементов, заменив последнюю строку на

    result = [np.empty(0, dtype=int)] * 3000   # Empty array, so OK to use same reference
    for i, j in enumerate(arr[locs]):
        result[j] = coords[i]
    return result

При желании вы можете фильтровать значения, которые находятся в указанном вами диапазоне c, который вы хотите (например, 0-2999).

Answer 2

Вы можете использовать логическое ИЛИ в numpy для одновременного прохождения всех этих условий равенства вместо одного за другим.

import numpy as np
conditions = False
for i in list:
  conditions = np.logical_or(conditions,array3d == i)

newlist = np.where(conditions)

Это позволяет numpy выполнять фильтрацию один раз вместо n проходов для каждого условие отдельно.

Еще один способ сделать это более компактно

np.where(np.isin(array3d, list))