Criação de uma nova coluna com base na condição if-elif-else


103

Eu tenho um DataFrame df:

    A    B
a   2    2 
b   3    1
c   1    3

Quero criar uma nova coluna com base nos seguintes critérios:

se linha A == B: 0

se linhaA > B: 1

se linha A < B: -1

então, dada a tabela acima, deveria ser:

    A    B    C
a   2    2    0
b   3    1    1
c   1    3   -1 

Para if elsecasos típicos que eu faço np.where(df.A > df.B, 1, -1), o pandas fornece uma sintaxe especial para resolver meu problema com uma etapa (sem a necessidade de criar 3 novas colunas e, em seguida, combinar o resultado)?


Você poderia apenas definir uma função e passar isso para applye definir axis=1funcionaria, não tenho certeza se posso pensar em uma operação que daria a você o que você deseja
EdChum

Sua solução implica criar 3 colunas e combiná-las em 1 coluna, ou você tem algo diferente em mente?
nozes

Você fica dizendo "criando 3 colunas", mas não tenho certeza do que está se referindo.
DSM

1
@DSM respondeu a esta pergunta, mas eu quis dizer algo como df['C']=df.apply(myFunc(row), axis=1)onde myFunc faz o que você quer, isso não envolve a criação de '3 colunas'
EdChum

1
Possível duplicação da criação condicional
Georgy

Respostas:


152

Para formalizar algumas das abordagens apresentadas acima:

Crie uma função que opere nas linhas de seu dataframe assim:

def f(row):
    if row['A'] == row['B']:
        val = 0
    elif row['A'] > row['B']:
        val = 1
    else:
        val = -1
    return val

Em seguida, aplique-o ao seu dataframe passando a axis=1opção:

In [1]: df['C'] = df.apply(f, axis=1)

In [2]: df
Out[2]:
   A  B  C
a  2  2  0
b  3  1  1
c  1  3 -1

Obviamente, isso não é vetorizado, portanto o desempenho pode não ser tão bom quando dimensionado para um grande número de registros. Ainda assim, acho que é muito mais legível. Especialmente vindo de um histórico SAS.

Editar

Aqui está a versão vetorial

df['C'] = np.where(
    df['A'] == df['B'], 0, np.where(
    df['A'] >  df['B'], 1, -1)) 

1
Obrigado, estou começando com pandas e isso foi muito útil +1
nutship

4
E se eu quiser passar outro parâmetro junto com a linha na função? Se eu fizer isso, dirá linha não definida ..
prashanth manohar

3
Você deve usar o argsparâmetro da .applyfunção: pandas.pydata.org/pandas-docs/stable/generated/…
Zelazny7

1
Sou um velho usuário do SAS aprendendo Python e há definitivamente uma curva de aprendizado! :-) Por exemplo, o código acima pode ser escrito em SAS como: data df; set df; if A=B then C=0; else if A>B then C=1; else C=-1; run;Muito elegante e simples.
RobertF

1
Uma resposta bem definida
Sahil Nagpal

54
df.loc[df['A'] == df['B'], 'C'] = 0
df.loc[df['A'] > df['B'], 'C'] = 1
df.loc[df['A'] < df['B'], 'C'] = -1

Fácil de resolver usando indexação. A primeira linha do código é assim: se coluna Afor igual a coluna B, crie e defina coluna Cigual a 0.


17

Para este relacionamento específico, você pode usar np.sign:

>>> df["C"] = np.sign(df.A - df.B)
>>> df
   A  B  C
a  2  2  0
b  3  1  1
c  1  3 -1

6

insira a descrição da imagem aqui

Digamos que um é o seu quadro de dados original e você deseja adicionar uma nova coluna 'antiga'

Se a idade for maior que 50, então consideramos como mais velho = sim, caso contrário, False

etapa 1: obter os índices de linhas com idade superior a 50

row_indexes=df[df['age']>=50].index

etapa 2: usando .loc, podemos atribuir um novo valor à coluna

df.loc[row_indexes,'elderly']="yes"

mesmo para idade abaixo de 50

row_indexes=df[df['age']<50].index

df[row_indexes,'elderly']="no"


1

Quando você tem várias if condições, numpy.selecté o caminho a seguir:

In [4102]: import numpy as np
In [4098]: conditions = [df.A.eq(df.B), df.A.gt(df.B), df.A.lt(df.B)]
In [4096]: choices = [0, 1, -1]

In [4100]: df['C'] = np.select(conditions, choices)

In [4101]: df
Out[4101]: 
   A  B  C
a  2  2  0
b  3  1  1
c  1  3 -1
Ao utilizar nosso site, você reconhece que leu e compreendeu nossa Política de Cookies e nossa Política de Privacidade.
Licensed under cc by-sa 3.0 with attribution required.