Получить индексы массива numpy, где значения смещены от себя на некоторое количество

Question

Если у меня есть два массива:

a = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
b = np.array([0, 1, 3, 4, 7, 8, 9])
n = 3

Как я могу без использования циклов найти индексы a, которые имеют значение, равное значениям b, смещенным на несколькочисло n?

Я думаю, что-то вроде этого будет работать, но я получаю предупреждение elementwise == comparison failed и пустой массив в результате:

np.where(a == b + n)

Вот цикл for, который выполняетчто я пытаюсь сделать:

for val in b:
    print(np.where(a == val + n))

Это выводит:

(array([3]),)
(array([4]),)
(array([6]),)
(array([7]),)
(array([], dtype=int64),)
(array([], dtype=int64),)
(array([], dtype=int64),)

Answer 1

Это должно сработать:

>>> a = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> b = np.array([0, 1, 3, 4, 7, 8, 9])
>>> n = 3
>>> np.where((a == b[:, np.newaxis] + n).any(axis=0))
(array([3, 4, 6, 7]),)

Идея состоит в том, чтобы транслировать a через b, поскольку вы хотите проверить все возможные значения b.В конце концов, это операция O(len(a) * len(b)), поэтому вам нужно создать двумерный массив размером len(a) * len(b) для векторизации.

Вещание достигается путем вставки другой оси в b, поэтомучто это вектор-столбец:

>>> b
array([0, 1, 3, 4, 7, 8, 9])
>>> b[:, np.newaxis]
array([[0],
       [1],
       [3],
       [4],
       [7],
       [8],
       [9]])

, так что теперь сравнение вернет двумерный массив:

>>> a == b[:, np.newaxis] + n
array([[False, False, False,  True, False, False, False, False, False, False],
       [False, False, False, False,  True, False, False, False, False, False],
       [False, False, False, False, False, False,  True, False, False, False],
       [False, False, False, False, False, False, False,  True, False, False],
       [False, False, False, False, False, False, False, False, False, False],
       [False, False, False, False, False, False, False, False, False, False],
       [False, False, False, False, False, False, False, False, False, False]])

Индексы i, j, где есть значение Trueзначит, что b[i] == a[j].Поскольку мы заботимся только о индексах a, которые соответствуют любому значению в b, мы просто ищем любое истинное значение по столбцам:

>>> (a == b[:, np.newaxis] + n).any(axis=0)
array([False, False, False,  True,  True, False,  True,  True, False, False])

Наконец, все, чтоздесь необходимо получить индексы, соответствующие значениям True:

>>> np.where((a == b[:, np.newaxis] + n).any(axis=0))
(array([3, 4, 6, 7]),)

---

В вашем примере ввода есть отсортированные массивы, что означает, что вы можете фактически выполнять линейную операцию в O(len(a) + len(b)), новы не сможете векторизовать его вообще.По сути, вы можете иметь индекс для каждого массива, увеличивая его по одному на b и продолжая увеличивать на a, пока не достигнете следующего значения.Конечно, это было бы менее эффективно в целом с массивами numpy, чем моё вышеупомянутое решение для «типичных» (т.е. не массивных) массивов.