CODEWITHSUNDEEP
pythonnumpy
Startiflette
Startiflette

Reputation: 111

Numpy : efficient matrix indexing

I have two numpy arrays :

Where k and n are integers.

The coefficients of B are indexes of the matrix A. i.e, for instance, B[0, 1, :] is an array [k, l] and we are sure that A[k , l] exists.

What I would like to do is to build a Matrix C of the same size as B, such that, for all i, j, C[i, j] = A[B[i, j, 0], B[i, j, 1]]

Is there an efficient way to do so ?

I have tried things like A[:, B], A[0, B], but it was unsuccesful. I could also do it with for loops but I think a implementation with numpy would be much faster.

For those who would like to try , I have prepared a little start-up code (with k=n=2 to test the methods :

import numpy as np

a = np.array([[
  [73,  -25],
  [97, -25],
  [73,  107],
  [97, 107]],

  [[81,  43],
  [86, 43],
  [81,  50],
  [86, 43]]
  ])

b = np.array(
  [[[0, 2],
    [0, 0],
    [0, 3],
    [0, 1],
    [1, 0],
    [1, 2],
    [1, 1],
    [1, 3],
    [0, 0],
    [0, 0],
    [0, 0],
    [0, 0]],

   [[0, 3],
     [0, 2],
     [0, 1],
     [0, 0],
     [1, 0],
     [1, 2],
     [1, 1],
     [1, 3],
     [0, 0],
     [0, 0],
     [0, 0],
     [0, 0]]]
  )

#the answer should be :
c = np.array(
  [[[73,  107],
    [73,  -25],
    [97, 107],
    [97, -25],
    [81,  43],
    [81,  50],
    [86, 43],
    [86, 50],
    [73,  -25],
    [73,  -25],
    [73,  -25],
    [73,  -25]],

   [[97, 107],
     [73,  107],
     [97, -25],
     [73,  -25],
     [81,  43],
     [81,  50],
     [86, 43],
     [86, 50],
     [73,  -25],
     [73,  -25],
     [73,  -25],
     [73,  -25]],
     ]
     )

Hope this is clear for you, Thanks in advance,

Upvotes: 1

Views: 58

Answers (2)

Quimby
Quimby

Reputation: 19123

You almost had it:

C=A[B[:,:,0],B[:,:,1]]

But you confused me for a minute since your test c has few incorrect values:

>>> a,b,c = ...
>>> C = a[b[:,:,0],b[:,:,1]]
>>> np.all([C[i,j]==a[b[i,j,0],b[i,j,1]] for i in range(2) for j in range(12)])
True
>>> C==c
array([[[ True,  True],
        [ True,  True],
        [ True,  True],
        [ True,  True],
        [ True,  True],
        [ True,  True],
        [ True, False],
        [ True, False],
        [ True,  True],
        [ True,  True],
        [ True,  True],
        [ True,  True]],

       [[ True,  True],
        [ True,  True],
        [ True,  True],
        [ True,  True],
        [ True,  True],
        [ True,  True],
        [ True, False],
        [ True, False],
        [ True,  True],
        [ True,  True],
        [ True,  True],
        [ True,  True]]])

Upvotes: 1

Dorian
Dorian

Reputation: 1488

a[b[...,0],b[...,1]]

should do the intended thing. You can do b[:, :, 0] instead of b[..., 0] but the latter one works for arrays with arbitrary dimensions, as long as you have your indexes at the last dimension.

Upvotes: 2

Related Questions