У меня есть список списков, отсортированных в порядке возрастания, аналогичный этому:

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

4
BaconBad 3 Янв 2018 в 15:09

2 ответа

Лучший ответ

Панды

Хотя это немного перебор, мы можем использовать для этого панды:

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, 2, None)]

Тогда получаем:

>>> [e for _, g in groupby(d, itemgetter(0)) for e in islice(g, 2, None)]
[[1, 3], [1, 4], [2, 3]]
6
Willem Van Onsem 3 Янв 2018 в 12:58

Вы также можете использовать 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]]
2
RoadRunner 3 Янв 2018 в 14:53