Numpyic way to take the first N rows and columns out of every M rows and columns from a square matrix

user_13860719 0
Viewed 118

I have a 20 x 20 square matrix. I want to take the first 2 rows and columns out of every 5 rows and columns, which means the output should be a 8 x 8 square matrix. This can be done in 2 consecutive steps as follows:

import numpy as np

m = 5 
n = 2 
A = np.arange(400).reshape(20,-1)
B = np.asarray([row for i, row in enumerate(A) if i % m < n])
C = np.asarray([col for j, col in enumerate(B.T) if j % m < n]).T

However, I am looking for efficiency. Is there a more Numpyic way to do this? I would prefer to do this in one step.

indexing python-3.x numpy arrays python
3 Answers

You can use np.ix_ to retain the elements whose row / column indices are less than 2 modulo 5:

import numpy as np

m = 5
n = 2
A = np.arange(400).reshape(20,-1)
mask = np.arange(20) % 5 < 2
result = A[np.ix_(mask, mask)]
print(result)

This outputs:

[[  0   1   5   6  10  11  15  16]
 [ 20  21  25  26  30  31  35  36]
 [100 101 105 106 110 111 115 116]
 [120 121 125 126 130 131 135 136]
 [200 201 205 206 210 211 215 216]
 [220 221 225 226 230 231 235 236]
 [300 301 305 306 310 311 315 316]
 [320 321 325 326 330 331 335 336]]
user_177698150

Very similar to accepted answered, but can just reference rows/column indices directly. Would be interested to see if benchmark is any different than using np.ix_() in accepted answer

Return Specific Row/Column by Numeric Indices

import numpy as np

m = 5 
n = 2 
A = np.arange(400).reshape(20,-1)
B = np.asarray([row for i, row in enumerate(A) if i % m < n])
C = np.asarray([col for j, col in enumerate(B.T) if j % m < n]).T

rowAndColIds = list(filter(lambda x: x % m < n,range(20)))
# print(rowAndColsIds)

result = A[:,rowAndColIds][rowAndColIds]
print (result)
user_46419280

You could use index broadcasting

i = (np.r_[:20:5][:, None] + np.r_[:2]).ravel()
A[i[:,None], i]

output:

array([[  0,   1,   5,   6,  10,  11,  15,  16],
       [ 20,  21,  25,  26,  30,  31,  35,  36],
       [100, 101, 105, 106, 110, 111, 115, 116],
       [120, 121, 125, 126, 130, 131, 135, 136],
       [200, 201, 205, 206, 210, 211, 215, 216],
       [220, 221, 225, 226, 230, 231, 235, 236],
       [300, 301, 305, 306, 310, 311, 315, 316],
       [320, 321, 325, 326, 330, 331, 335, 336]])
user_149232270
Related