CODEWITHSUNDEEP
pythonnumpy
blahblah
blahblah

Reputation: 1251

Removing `n` elements from each row in `2D` numpy where `n` varies row by row

How would one remove n elements from 2D numpy array where n varies in each row?

For example:

# input array 
[[1,2,3],
 [3,1,2],
 [1,2,3],
 [4,5,2],
 [5,6,7]]

with

# n elements to remove from each row
[0, 2, 1, 2, 1]

would result in:

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

Do note that the result does not need to be a numpy array(and won't be as Michael noticed in the comments), just an arbitrary Python list of lists.

Upvotes: 0

Views: 96

Answers (3)

fountainhead
fountainhead

Reputation: 3722

If arr is your numpy array, rem is a list of removal-counts, a simple solution could be:

res = [arr[i][:3-rem[i]].tolist() for i in range(len(rem))]

Output:

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

Upvotes: 0

de1
de1

Reputation: 3124

As has been pointed out in the comments, the result wouldn't be a numpy array.

Given the example data from your question:

>>> a = [[1,2,3],
...  [3,1,2],
...  [1,2,3],
...  [4,5,2],
...  [5,6,7]]
>>> remove_n_tail_list = [0, 2, 1, 2, 1]

You could for example use a list comprehension to get the desired result:

>>> [row[:len(row) - remove_n_tail] for row, remove_n_tail in zip(a, remove_n_tail_list)]
[[1, 2, 3], [3], [1, 2], [4], [5, 6]]

In that solution, row[:len(row) - remove_n_tail] is selecting the values up to the length of the row (len(row)), minus the number of values you want to remove from the end (remove_n_tail).

There are various methods to achieve similar results. You might find the Most efficient way to map function over numpy array question interesting.

Upvotes: 2

ZealousWeb
ZealousWeb

Reputation: 1751

Herewith I have come up with a solution. Give it a try:

arr = [[1,2,3],
 [3,1,2],
 [1,2,3],
 [4,5,2],
 [5,6,7]]
new_arr =[]
lst = [0, 2, 1, 2, 1]
count=0
for i in arr:
    remove = lst[count]
    count = count+1
    temp = i[:len(i)-remove]
    new_arr.append(temp)
print(new_arr)

Upvotes: 2

Related Questions