过滤列表中的前两个匹配元素

Posted 2023-02-22

【中文标题】过滤列表中的前两个匹配元素

【英文标题】：Filtering the two first matching elements in a list

【发布时间】：2018-06-13 02:42:00

【问题描述】：

我有一个按升序排列的列表，类似于这个：

input = [[1,1],[1,2],[1,3],[1,4],[2,1],[2,2],[2,3],[3,1],[6,1],[6,2]]

我想过滤这个列表，以便新列表只包含前两个（或唯一）在位置 0 处具有匹配整数的元素，如下所示：

output = [[1,1],[1,2],[2,1],[2,2],[3,1],[6,1],[6,2]]

如果剩余元素（不符合条件的元素）将保留在输入列表中，而匹配的元素将单独存储，那将是理想的。

我该怎么做？

提前谢谢你！

编辑：索引 1 上的元素几乎可以是任何整数，例如[[1,6],[1,7],[1,8],[2,1],[2,2]]

【问题讨论】：

@WillemVanOnsem 啊，好吧。

OP，你做了什么？这可以使用itertools.groupby 在一行中完成。

@cᴏʟᴅsᴘᴇᴇᴅ 正确，如果它们重复多次，它们仍保留在输入列表中，在这种情况下，过滤后为[[1,3],[1,4],[2,3]]。我说得通吗？

【参考方案1】：

熊猫

虽然这有点矫枉过正，但我​​们可以使用 pandas：

import pandas as pd

pd.DataFrame(d).groupby(0).head(2).values.tolist()

带有d的原始列表。然后产生：

>>> pd.DataFrame(d).groupby(0).head(2).values.tolist()
[[1, 1], [1, 2], [2, 1], [2, 2], [3, 1], [6, 1], [6, 2]]

请注意，这将返回列表的副本，而不是原始列表。此外，所有行都应具有相同数量的项目。

Itertools groupby 和 islice

如果列表是按字典顺序排列的，那么我们可以使用itertools.groupby：

from operator import itemgetter
from itertools import groupby, islice

[e for _, g in groupby(d, itemgetter(0)) for e in islice(g, 2)]

这又产生了：

>>> [e for _, g in groupby(d, itemgetter(0)) for e in islice(g, 2)]
[[1, 1], [1, 2], [2, 1], [2, 2], [3, 1], [6, 1], [6, 2]]

它也更加灵活，因为我们将引用复制到列表中，并且所有列表可以有不同数量的元素（这里至少一个）。

编辑

可以通过让islice 以相反的方式工作来获得其余值：保留除前两个之外的所有内容：

[e for _, g in groupby(d, itemgetter(0)) for e in islice(g<b>, 2, None</b>)]

然后我们得到：

>>> [e for _, g in groupby(d, itemgetter(0)) for e in islice(g, 2, None)]
[[1, 3], [1, 4], [2, 3]]

【讨论】：

itertools 解决方案可以完成这项工作。谢谢！在 for 循环中使用 _ 符号有什么特别的原因吗？

@BaconBad：是的，它是解包元组，通常你写_ 来表示“无关变量”。完成这项工作所需的变量，但我们对其包含的值不感兴趣。

明白了！最后一件事：有没有一种方法可以方便地检索遗漏的值（原始列表减去新过滤的值），而无需在之后对新旧列表进行手动比较？

@BaconBad：我将对d 列表执行第二次遍历。它需要线性时间，因此速度相当快。

【参考方案2】：

您还可以使用collections.defaultdict 按第一个索引对子列表进行分组：

from collections import defaultdict
from pprint import pprint

input_lst = [[1,1],[1,2],[1,3],[1,4],[2,1],[2,2],[2,3],[3,1],[6,1],[6,2]]

groups = defaultdict(list)
for lst in input_lst:
    key = lst[0]
    groups[key].append(lst)

pprint(groups)

这给出了这个分组字典：

defaultdict(<class 'list'>,
        1: [[1, 1], [1, 2], [1, 3], [1, 4]],
         2: [[2, 1], [2, 2], [2, 3]],
         3: [[3, 1]],
         6: [[6, 1], [6, 2]])

然后您可以从每个键中获取前两个 [:2] 值，并确保结果被展平并最终排序：

from itertools import chain

result = sorted(chain.from_iterable(x[:2] for x in groups.values()))

print(result)

哪些输出：

[[1, 1], [1, 2], [2, 1], [2, 2], [3, 1], [6, 1], [6, 2]]