Como implementar 'in' e 'not in' para um DataFrame de pandas?
O Pandas oferece dois métodos: Series.isin
e DataFrame.isin
para Series e DataFrames, respectivamente.
Filtrar DataFrame com base em uma coluna (também se aplica a séries)
O cenário mais comum é aplicar uma isin
condição em uma coluna específica para filtrar linhas em um DataFrame.
df = pd.DataFrame({'countries': ['US', 'UK', 'Germany', np.nan, 'China']})
df
countries
0 US
1 UK
2 Germany
3 China
c1 = ['UK', 'China'] # list
c2 = {'Germany'} # set
c3 = pd.Series(['China', 'US']) # Series
c4 = np.array(['US', 'UK']) # array
Series.isin
aceita vários tipos como entradas. A seguir, são apresentadas todas as formas válidas de obter o que você deseja:
df['countries'].isin(c1)
0 False
1 True
2 False
3 False
4 True
Name: countries, dtype: bool
# `in` operation
df[df['countries'].isin(c1)]
countries
1 UK
4 China
# `not in` operation
df[~df['countries'].isin(c1)]
countries
0 US
2 Germany
3 NaN
# Filter with `set` (tuples work too)
df[df['countries'].isin(c2)]
countries
2 Germany
# Filter with another Series
df[df['countries'].isin(c3)]
countries
0 US
4 China
# Filter with array
df[df['countries'].isin(c4)]
countries
0 US
1 UK
Filtrar em MUITAS colunas
Às vezes, convém aplicar uma verificação de associação "com" com alguns termos de pesquisa em várias colunas,
df2 = pd.DataFrame({
'A': ['x', 'y', 'z', 'q'], 'B': ['w', 'a', np.nan, 'x'], 'C': np.arange(4)})
df2
A B C
0 x w 0
1 y a 1
2 z NaN 2
3 q x 3
c1 = ['x', 'w', 'p']
Para aplicar a isin
condição às duas colunas "A" e "B", use DataFrame.isin
:
df2[['A', 'B']].isin(c1)
A B
0 True True
1 False False
2 False False
3 False True
A partir disso, para manter as linhas onde há pelo menos uma colunaTrue
, podemos usar any
ao longo do primeiro eixo:
df2[['A', 'B']].isin(c1).any(axis=1)
0 True
1 False
2 False
3 True
dtype: bool
df2[df2[['A', 'B']].isin(c1).any(axis=1)]
A B C
0 x w 0
3 q x 3
Observe que, se você deseja pesquisar todas as colunas, omite a etapa de seleção de colunas e faça
df2.isin(c1).any(axis=1)
Da mesma forma, para manter as linhas onde estão TODAS as colunasTrue
, use all
da mesma maneira que antes.
df2[df2[['A', 'B']].isin(c1).all(axis=1)]
A B C
0 x w 0
Notável Menções: numpy.isin
,query
, compreensões lista (dados de cadeia)
Além dos métodos descritos acima, você também pode usar o equivalente numpy: numpy.isin
.
# `in` operation
df[np.isin(df['countries'], c1)]
countries
1 UK
4 China
# `not in` operation
df[np.isin(df['countries'], c1, invert=True)]
countries
0 US
2 Germany
3 NaN
Por que vale a pena considerar? As funções do NumPy geralmente são um pouco mais rápidas que seus equivalentes de pandas, devido à menor sobrecarga. Como essa é uma operação elementar que não depende do alinhamento do índice, há muito poucas situações em que esse método não é um substituto apropriado para os pandas.isin
.
As rotinas do Pandas geralmente são iterativas ao trabalhar com strings, porque é difícil vetorizar operações de strings. Há muitas evidências para sugerir que a compreensão da lista será mais rápida aqui. . Recorremos a um in
cheque agora.
c1_set = set(c1) # Using `in` with `sets` is a constant time operation...
# This doesn't matter for pandas because the implementation differs.
# `in` operation
df[[x in c1_set for x in df['countries']]]
countries
1 UK
4 China
# `not in` operation
df[[x not in c1_set for x in df['countries']]]
countries
0 US
2 Germany
3 NaN
No entanto, é muito mais difícil especificar, portanto, não o use, a menos que você saiba o que está fazendo.
Por fim, há também o DataFrame.query
que foi abordado nesta resposta . numexpr FTW!