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