Is there a way to do array-based indexing in NumPy so that indices corresponding to all possible combinations of indexing array elements are selected?

Question

arr = np.zeros([5,5])
idx_arr = np.array([1,3,4])

arr[idx_arr, idx_arr] += 1
print(arr)

Above gives the following output:

array([[0., 0., 0., 0., 0.],
       [0., 1., 0., 0., 0.],
       [0., 0., 0., 0., 0.],
       [0., 0., 0., 1., 0.],
       [0., 0., 0., 0., 1.]])

I'm looking for the following output:

array([[0., 0., 0., 0., 0.],
       [0., 1., 0., 1., 1.],
       [0., 0., 0., 0., 0.],
       [0., 1., 0., 1., 1.],
       [0., 1., 0., 1., 1.]])

In R, this form of indexing would be equivalent to:

df[c(1,2,3),c(1,2,3)]

Which would produce a 3 x 3 data frame or matrix sub-section.

Any help would be greatly appreciated!

Broadcast here. `arr[idx_arr[:, None], idx_arr] +=1` – user3483203 Feb 10 '20 at 16:13 — user3483203, Feb 10 '20 at 16:13

score 2 · Answer 1 · answered Feb 10 '20 at 16:14

You can obtain this result by adding a dimension to one of the indexing arrays so the smaller is broadcast across the larger:

arr[idx_arr[:,None], idx_arr] += 1
print(arr)

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

Or also using np.ix_:

arr[np.ix_(idx_arr, idx_arr)] += 1

1 Answers1