CODEWITHSUNDEEP
pythonnumpyindexing
rpb
rpb

Reputation: 3299

Faster approach extracting lower triangle of multidimensional array using Numpy?

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?

Upvotes: 2

Views: 122

Answers (3)

Mad Physicist
Mad Physicist

Reputation: 114310

Since we are iteratively improving the indexing expressions:

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

Upvotes: 2

Daniel F
Daniel F

Reputation: 14399

You can also do:

arr[lw_up_pair][:, 0]

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

Upvotes: 2

Dani Mesejo
Dani Mesejo

Reputation: 61910

Just do:

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

Output

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

Upvotes: 3

Related Questions