CODEWITHSUNDEEP
pythonpython-3.xnumpynumpy-ndarray
Clement Attlee
Clement Attlee

Reputation: 733

Select elements from numpy array based on a list of lengths

Suppose I have a numpy array like this:

arr = np.array([[1,1,1,1], [2,2,2,2], [3,3,3,3], [4,4,4,4])

I also have a list that determines the intended lengths:

lens = [1,2,3,4]

Is there an elegant and Pythonic way to return a new array, with each corresponding element selected using the lens variable?

The output should be:

[[1], [2,2],[3,3,3], [4,4,4,4]]

Upvotes: 1

Views: 477

Answers (3)

RAGHHURAAMM
RAGHHURAAMM

Reputation: 1099

Assuming arr is equal in length as lens: 

[arr[i][:v].tolist() for i,v in enumerate(lens)]

Output: [[1], [2, 2], [3, 3, 3], [4, 4, 4, 4]]

Upvotes: 0

Chris
Chris

Reputation: 29752

If each Numpy arr list element size is less than 5 this works.

Use zip:

[a[:i].tolist() for a, i in zip(arr, lens)]

Output:

[[1], [2, 2], [3, 3, 3], [4, 4, 4, 4]]

Upvotes: 5

Bouwun Tsui
Bouwun Tsui

Reputation: 1

This could be a possible solution: 

res = []
for a,i in zip(arr, lens):
    res.append(a[:i].tolist())
print(res)

Upvotes: 0

Related Questions