如何在保持筛选值为零的情况下筛选Python列表

与

输入=[0,0,5,9,0,4,10,3,0]

作为列表 我需要一个输出，它将是输入中的两个最高值，同时将其他列表元素设置为零

输出=[0,0,0,9,0,0,10,0,0]

我得到的最接近的结果是：

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

---

您能帮忙吗？

您可以排序，找到两个最大的值，然后使用列表：

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]

输出：

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

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

---

可以通过单个列表遍历实现这一点，使算法O（n）：

• 首先通过一次遍历找到两个最高值

• 然后创建一个零列表，并添加找到的最大值

代码 请注意，如果列表中的最大值出现两次以上，则仅显示最左侧的两个值

---

作为旁注，不要在代码中使用诸如

input

之类的名称，因为这会掩盖相同名称的内置函数。

不像Ajax的解决方案那么漂亮，而是一个

O（n）

解决方案和一个更具动态性的：

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

---

纯numpy方法：

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)

输出：

[ 0  0  0  9  0  0 10  0  0]

---

这里是另一个基于

numpy

的解决方案，它可以避免对整个阵列进行排序，这需要

O（nlogn）

时间

import numpy as np

arr = np.array([0,0,5,9,0,4,10,3,0])
arr[np.argpartition(arr,-2)[:-2]] = 0

如果要创建新阵列作为输出：

result = np.zeros_like(arr)
idx = np.argpartition(arr,-2)[-2:]
result[idx] = arr[idx]

相应的Python本机解决方案是使用

heap.nlargest

，这也避免了对整个数组进行排序

import heapq

arr = [0,0,5,9,0,4,10,3,0]
l = len(arr)
idx1, idx2 = heapq.nlargest(2, range(l), key=arr.__getitem__)
result = [0] * l
result[idx1] = arr[idx1]
result[idx2] = arr[idx2]

---

如果最大值出现两次以上，可能会重复什么？虽然它很整洁，但这是O（n log（n）），但可以在O（n）中使用单个traversal@Ajax1234@Olivier对于O（n）解决方案，您可以使用

heapq.nlargest

而不是

排序

。

import heapq

arr = [0,0,5,9,0,4,10,3,0]
l = len(arr)
idx1, idx2 = heapq.nlargest(2, range(l), key=arr.__getitem__)
result = [0] * l
result[idx1] = arr[idx1]
result[idx2] = arr[idx2]