3

I have a list of lists of indices for an numpy array, but do not quite come to the wanted result when using them.

n = 3
a = np.array([[8, 1, 6],
              [3, 5, 7],
              [4, 9, 2]])
np.random.seed(7)
idx = np.random.choice(np.arange(n), size=(n, n-1))
# array([[0, 1],
#        [2, 0],
#        [1, 2]])

In this case I want:

  • element 0 and 1 of row 0
  • element 2 and 0 of row 1
  • element 1 and 2 of row 2

My list has n sublists and all of those lists have the same length. I want that each sublist is only used once and not for all axis.

# Wanted result
# b = array[[8, 1],
#           [7, 3],
#           [9, 2]])

I can achieve this but it seems rather cumbersome with a lot of repeating and reshaping.

# Possibility 1
b = a[:, idx]
# array([[[8, 1],   | [[3, 5],   |  [[4, 9],
#         [6, 8],   |  [7, 3],   |   [2, 4],
#         [1, 6]],  |  [5, 7]],  |   [9, 2]])
b = b[np.arange(n), np.arange(n), :]

# Possibility 2
b = a[np.repeat(range(n), n-1), idx.ravel()]
# array([8, 1, 7, 3, 9, 2])
b = b.reshape(n, n-1)

Are there easier ways?

scleronomic
  • 4,392
  • 1
  • 13
  • 43

1 Answers1

2

You can use np.take_along_axis here:

np.take_along_axis(a, idx, 1)

array([[8, 1],
       [7, 3],
       [9, 2]])

Or using broadcasting:

a[np.arange(a.shape[0])[:,None], idx]

array([[8, 1],
       [7, 3],
       [9, 2]])

Note that your using integer array indexing here, you need to specify over which axis and rows you want to index using idx.

yatu
  • 86,083
  • 12
  • 84
  • 139