从 `2D` numpy 的每一行中删除 `n` 元素，其中 `n` 逐行变化

Question

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.

Answer 1

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)   

  

Answer 2

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.

Answer 3

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]]