How to filter Python list while keeping filtered values zero

Question

With

input = [0,0,5,9,0,4,10,3,0]

as list I need an output, which is going to be two highest values in input while setting other list elements to zero.

output = [0,0,0,9,0,0,10,0,0]

The closest I got:

from itertools import compress
import numpy as np
import operator
input= [0,0,5,9,0,4,10,3,0]
top_2_idx = np.argsort(test)[-2:]
test[top_2_idx[0]]
test[top_2_idx[1]]

Can you please help?

Answer 1

You can sort, find the two largest values, and then use a list comprehension:

input = [0,0,5,9,0,4,10,3,0]
*_, c1, c2 = sorted(input)
result = [0 if i not in {c1, c2} else i for i in input]

Output:

[0, 0, 0, 9, 0, 0, 10, 0, 0]

Answer 2

Not as pretty as Ajax's solution but a O(n) solution and a little more dynamic:

from collections import deque

def zero_non_max(lst, keep_top_n):
    """
    Returns a list with all numbers zeroed out
    except the keep_top_n.
    >>> zero_non_max([0, 0, 5, 9, 0, 4, 10, 3, 0], 3)
    >>> [0, 0, 5, 9, 0, 0, 10, 0, 0]
    """

    lst = lst.copy()

    top_n = deque(maxlen=keep_top_n)

    for index, x in enumerate(lst):
        if len(top_n) < top_n.maxlen or x > top_n[-1][0]:
            top_n.append((x, index))
        lst[index] = 0

    for val, index in top_n:
        lst[index] = val

    return lst

lst = [0, 0, 5, 9, 0, 4, 10, 3, 0]
print(zero_non_max(lst, 2))

Output:

[0, 0, 0, 9, 0, 0, 10, 0, 0]

Answer 3

Pure numpy approach:

import numpy as np

arr = np.array([0, 0, 5, 9, 0, 4, 10, 3, 0])
top_2_idx = np.argsort(arr)[-2:]
np.put(arr, np.argwhere(~np.isin(arr, arr[top_2_idx])), 0)
print(arr)

The output:

[ 0  0  0  9  0  0 10  0  0]

Numpy.put