CODEWITHSUNDEEP
pythonnumpy
Cron Merdek
Cron Merdek

Reputation: 1124

Select numpy array of different shape by index and write it back

I looking for indexing strategy how I can select sub-array(right-side) from original array(left-side) and then replace original array with sub-array(bottom). There are indexes to select for each row, e.g. for first row they are [1,3,0], and for last are [4,6,2]. Full selection matrix for the illustration below looks like this:

[[1,3,0],
 [2,3,2],
 [4,0,6],
 [1,4,0],
 [4,6,2]]

So far I can make only looping through rows of original array and replace values at column indexes. Is there solution without a loop?

         +--------------------------------------+
         |                                      |
         |           +-----------------------+  |
         |           |                       |  |
         |   +----------------------------+  |  |
         |   |       |                    |  |  |
         |   |       |                    |  |  |
         |   |       |                    |  |  |
         |   |       |                    |  |  |
         +   +       +                    v  v  v
array([[ 0,  1,  2,  3,  4,  5,  6],      1, 3, 0,
       [ 7,  8,  9, 10, 11, 12, 13],      9,10,9,
       !14, 15, 16, 17,!18, 19,!20],     18,14,20,
       [21, 22, 23, 24, 25, 26, 27],     22,25,21,
       [28, 29, 30, 31, 32, 33, 34]])    32,34,30
                +        +       +        ^  ^  ^
                |        |       |        |  |  |
                |        +----------------+  |  |
                |                |           |  |
                |                +-----------+  |
                |                               |
                +-------------------------------+

Final array is original array with replaced last indices from sub-array( which is selected from original array.

                       +------------+
array([[ 0,  1,  2,  3,| 1,  3,  0],|
       [ 7,  8,  9, 10,| 9, 10,  9],|
       [14, 15, 16, 17,|18, 14, 20],|
       [21, 22, 23, 24,|22, 25, 21],|
       [28, 29, 30, 31,|32, 34, 30]]|
                       +------------+

Solution with loop:

selection = np.array([[1,3,0],
                     [2,3,2],
                     [4,0,6],
                     [1,4,0],
                     [4,6,2]])
y = np.arange(35).reshape(5,7)
s = y.shape[1] - selection.shape[1]
e = y.shape[1]
for i in range(0, y.shape[0]):
    y[i, s:e] = y[i, selection[i]]

>>> y
array([[ 0,  1,  2,  3,  1,  3,  0],
       [ 7,  8,  9, 10,  9, 10,  9],
       [14, 15, 16, 17, 18, 14, 20],
       [21, 22, 23, 24, 22, 25, 21],
       [28, 29, 30, 31, 32, 34, 30]])

Upvotes: 3

Views: 473

Answers (2)

Paul Panzer
Paul Panzer

Reputation: 53029

Use broaodcasting and advanced indexing:

>>> y[:, 4:] = y[np.arange(5)[:, None], selection]
>>> y
array([[ 0,  1,  2,  3,  1,  3,  0],
       [ 7,  8,  9, 10,  9, 10,  9],
       [14, 15, 16, 17, 18, 14, 20],
       [21, 22, 23, 24, 22, 25, 21],
       [28, 29, 30, 31, 32, 34, 30]])

Upvotes: 2

javidcf
javidcf

Reputation: 59711

You can do it by creating the indices for the rows:

import numpy as np

a = np.array([[ 0,  1,  2,  3,  4,  5,  6],
              [ 7,  8,  9, 10, 11, 12, 13],
              [14, 15, 16, 17, 18, 19, 20],
              [21, 22, 23, 24, 25, 26, 27],
              [28, 29, 30, 31, 32, 33, 34]])

idx = np.array([[1,3,0],
                [2,3,2],
                [4,0,6],
                [1,4,0],
                [4,6,2]])

rows = np.tile(np.arange(len(idx)), [idx.shape[1], 1]).T
# This array looks like this:
# [[0, 0, 0],
#  [1, 1, 1],
#  [2, 2, 2],
#  [3, 3, 3],
#  [4, 4, 4]]

a[:, -idx.shape[1]:] = a[rows, idx]
print(a)

Output:

array([[ 0,  1,  2,  3,  1,  3,  0],
       [ 7,  8,  9, 10,  9, 10,  9],
       [14, 15, 16, 17, 18, 14, 20],
       [21, 22, 23, 24, 22, 25, 21],
       [28, 29, 30, 31, 32, 34, 30]])

Upvotes: 2

Related Questions