Векторный поиск элемента в значении - списке?

Question

Как получить список планет по критериям отбора?

DataFrame.isin для поиска значения в списке как здесь не подходит.
Очень желательно векторный способ Panadas, потому что тесты на большой базе данных.

import pandas as pd

df = pd.DataFrame({
    'name': ['Earth', 'Jupiter', 'Mars'],
    'mass_to_earth': [1.0, 317.8, 0.1],
    'kind': ['planet', 'planet', 'planet'],
    'colors': [['red','blue', 'green'], ['blue', 'green'], ['red','blue']]
})

print(df)

choice = df['name'][
    (df['mass_to_earth'] > 0.2) &
    (df['kind'].str.contains('planet', case=True, regex=False, na=False)) &
    (df['colors'].isin(['green'])) #это не работает как итератор
].tolist()

print(type(choice),choice)

И как вариант поиск в set пожалуйста для теста скорости:

'colors': [{'red','blue', 'green'}, {'blue', 'green'}, {'red','blue'}]

Answer 1

Вариант 1 - преобразовать список в CSV строку и искать при помощи регулярного выражения:

In [169]: df.loc[df["colors"].str.join(sep="|").str.contains(r"\bgreen\b")]
Out[169]:
      name  mass_to_earth    kind              colors
0    Earth            1.0  planet  [red, blue, green]
1  Jupiter          317.8  planet       [blue, green]

PS В микроязыке RegEx \b обозначает границу слова

Вариант 2: сначала преобразуем списки в множества (операция "element in set" - имеет асимптотическую сложность O(1), а для списков это O(n)), потом поиск "в лоб":

df["colors"] = df["colors"].map(set)

In [176]: df.loc[["green" in x for x in df["colors"]]]
Out[176]:
      name  mass_to_earth    kind              colors
0    Earth            1.0  planet  {red, green, blue}
1  Jupiter          317.8  planet       {green, blue}

PS на первый взгляд этот вариант может показаться медленным из-за использования цикла, но на практике, такой вариант часто оказывается быстрее остальных для строковых типов данных. Нужно тестировать и сравнивать время выполнения на реальных данных...

Answer 2

df['colors'].isin(['green']) # ошибка

Вот такой вариант условия нормально работает:

df['colors'].apply(lambda x: 'green' in x)

Вам по сути нужен contains для списков/множеств, его можно с помощью lambda так вот сделать.

Answer 3

Поскольку имела значение еще и скорость выполнения, была подобрана самая подходящая функция - Series.map в паре с sets.issubset(), она на тестах отрабатывала стабильно повторяемо на 25-30% быстрее ближайшего конкурента Series.apply, итого:

import pandas as pd

df = pd.DataFrame({
    'name': ['Earth', 'Jupiter', 'Mars'],
    'mass_to_earth': [1.0, 317.8, 0.1],
    'kind': ['planet', 'planet', 'planet'],
    'colors': [{'red','blue', 'green'}, {'blue', 'green'}, {'red','blue'}] # set или frozenset 
})

print(df)

choice = df['name'][
    (df['mass_to_earth'] > 0.2) &
    (df['kind']=='planet') &
    (df['colors'].map({'green'}.issubset))
].tolist()

print(choice)

UPD: Проверил, df['colors'].apply({'green'}.issubset) то же по скорости, вся фишка в sets.issubset().

Оставляю принятым ответ CrazyElf, как первый в Pandas style )

БЛОГ НА HUSL

Векторный поиск элемента в значении - списке?

Ответы (3 шт):