Эффективный способ получить подмножество индексов в NumPy

Question

У меня есть следующие индексы, которые вы получили бы от np.where(...):

coords = (
  np.asarray([0 0 0 1 1 1 1 1 2 2 2 3 3 3 3 4 4 4 5 5 5 5 5 6 6 6]),
  np.asarray([2 2 8 2 2 4 4 6 2 2 6 2 2 4 6 2 2 6 2 2 4 4 6 2 2 6]),
  np.asarray([0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]),
  np.asarray([0 1 0 0 1 0 1 1 0 1 1 0 1 1 1 0 1 1 0 1 0 1 1 0 1 1])
)

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

index = tuple(
  np.asarray([0 0 1 1 1 1 2 2 2 3 3 3 3 4 4 4 5 5 5 5 5 6 6 6]),
  np.asarray([2 8 2 4 4 6 2 2 6 2 2 4 6 2 2 6 2 2 4 4 6 2 2 6]),
  np.asarray([0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]),
  np.asarray([0 0 1 0 1 1 0 1 1 0 1 1 1 0 1 1 0 1 0 1 1 0 1 1])
)

Так, например, ords [0] выбрано, потому что оно в индексе (в позиции 0), но coords[1] не выбрано, потому что оно недоступно в index.

Я могу легко вычислить маску с помощью [x in zip(*index) for x in zip(*coords)] (преобразовано из bool в int для лучшей читаемости):

[1 0 1 0 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1]

но это не очень эффективно для больших массивов. Есть ли более «основанный на numpy» способ, которым можно вычислить маску?

Answer 1

Вы можете использовать от np.ravel_multi_index до , сжимая столбцы в уникальные числа , которые легче обрабатывать:

cmx = *map(np.max, coords),
imx = *map(np.max, index),
shape = np.maximum(cmx, imx) + 1

ct = np.ravel_multi_index(coords, shape)
it = np.ravel_multi_index(index, shape)

it.sort()

result = ct == it[it.searchsorted(ct)]
print(result.view(np.int8))

Печать:

[1 0 1 0 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1]

Answer 2

Не уверен насчет эффективности, но, учитывая, что вы в основном сравниваете пары координат, вы можете использовать scipy функции расстояния. Что-то вместе:

from scipy.spatial.distance import cdist

c = np.stack(coords).T
i = np.stack(index).T

d = cdist(c, i)

In [113]: np.any(d == 0, axis=1).astype(int)
Out[113]: 
array([1, 0, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
       1, 1, 1, 1])

По умолчанию он использует норму L2, возможно, вы могли бы сделать это немного быстрее с помощью более простой функции расстояния, например ::

d = cdist(c,i, lambda u, v: np.all(np.equal(u,v)))
np.any(d != 0, axis=1).astype(int)