Блог

У меня есть большой список из ~ 10000 списков из трех элементов, и мне нужно найти в нем значение. Мне приходится проделывать эту операцию много раз — над одними и теми же данными — только с другим значением.

Например, если список:

 l = [[1, 4, 5], [3, 7, 1], [2, 6, 8]]
  

и значение, которое я ищу 1 , я могу искать списки, содержащие значение, выполнив:

 [sublist for sublist in l if value in sublist]
  

Для этого примера ожидаемый результат будет:

 [[1, 4, 5], [3, 7, 1]]
  

Есть ли более быстрый способ сделать это?

Предполагая, что значения в подсписках являются хешируемыми, вы можете создать таблицу поиска — dict, где ключами являются значения в подсписках, а значениями dict — коллекция индексов подсписков во внешнем списке.

 from collections import defaultdict
spam = [[1, 4, 5], [3, 7, 1], [2, 6, 8]]
lookup = defaultdict(list)
for idx, item in enumerate(spam):
    for key in item:
        lookup[key].append(idx)

def get_sublists(array, lookup_table, value):
    return [array[idx] for idx in lookup_table[value]]


print(get_sublists(spam, lookup, 1))
print(get_sublists(spam, lookup, 7))
  

вывод

 [[1, 4, 5], [3, 7, 1]]
[[3, 7, 1]]
  

РЕДАКТИРОВАТЬ: некоторое сравнение с ответом @skarit

 from collections import defaultdict
from random import sample

spam = [list(sample(range(10), 3)) for _ in range(10000)]

lookup = defaultdict(list)
for idx, item in enumerate(spam):
    for key in item:
        lookup[key].append(idx)


def get_sublists(array, lookup_table, value):
    return [array[idx] for idx in lookup_table[value]]

import pandas as pd
import timeit


df = pd.DataFrame(spam)

print(timeit.timeit('[sublist for sublist in spam if 1 in sublist]', setup="from __main__ import spam", number=10000))
print(timeit.timeit('df = df[(df == 1).any(1)]', setup="from __main__ import df", number=10000))
print(timeit.timeit('get_sublists(spam, lookup, 1)', setup="from __main__ import get_sublists, spam, lookup", number=10000))
  

вывод

 10.568637077001767
10.028736718999426
2.4226377830018464
  

EDIT2: следуя комментарию @HeapOverfow, вот версия, в которой таблица поиска будет хранить соответствующие подсписки вместо их индексов. В зависимости от размера списка, следите за объемом памяти. Приведенный ниже фрагмент также отражает комментарий OP под ответом, чтобы исключить значение поиска из подсписков.

 from collections import defaultdict
from random import sample

spam = [list(sample(range(10), 3)) for _ in range(1000)]

lookup = defaultdict(list)
for item in spam:
    for key in item:
        lookup[key].append([num for num in item if num != key])


def get_sublists(array, lookup_table, value):
    return lookup_table[value]
  

Попробуйте это:

 import pandas as pd
value = 1 # The value you want to search

l = [[1, 4, 5], [3, 7, 1], [2, 6, 8]]
df = pd.DataFrame(l)
df = df[(df == value).any(1)]
  

Теперь этот фрейм данных будет содержать только «списки», которые имеют значение в виде строк