Faster approach extracting lower triangle of multidimensional array using Numpy?

user_6446053 0
Viewed 62

Assume a 4D array of shape 4,4,1,1. The objective is to extract the lower-triangle of an (n, m) array.

Currently, the following code which rely on a for loop, does the job.

import numpy as np
lw_up_pair = np.tril_indices(4, -1)
arr=np.zeros((4,4,1,1))

arr[1,:1,:,0]=1
arr[2,:2,0,0]=2
arr[3,:3,0,0]=3

ext=[arr [x, i, :, 0] for x,i in zip(lw_up_pair [0],lw_up_pair [1])] # expected output

However, I wonder whether there is more faster or numpy build-in approach that can achieve the similar result?

indexing numpy python
3 Answers

Just do:

rows, cols = lw_up_pair
ext = arr[rows, cols, :, 0]
print(ext)

Output

[[1.]
 [2.]
 [2.]
 [3.]
 [3.]
 [3.]]
user_40015920

You can also do:

arr[lw_up_pair][:, 0]

array([[1.],
       [2.],
       [2.],
       [3.],
       [3.],
       [3.]])
user_44277770

Since we are iteratively improving the indexing expressions:

arr[(*np.tril_indices(4, -1), slice(None), 0)]
user_29887300
Related