Search code examples
pythonpandasdataframeif-statementvariables

Is there a function or way to create a new variable based on the result of multiple if statement in python?


I have a problem of creating new variable based on multiple if statements. Programming language - python.

#for demonstration
import pandas as pd

example = {
  "blue": [1, 1, 0, 0, 1],
  "red": [1, 1, 0, 1, 1],
    "green": [0,0, 2, 0, 1]
}

#load into df:
example = pd.DataFrame(example)

print(example) 

New variable should be created based on this condition: if blue or/and red=1 and green = 0.

so, I come up with this solution, which only allows me to print the result

if example["blue"].iloc[0]==1 or example["red"].iloc[0]==1:
    if example["blue"].iloc[0]==1 and example["red"].iloc[0]==1:
        if example["green"].iloc[0]==0:
            print("Flower A")
        else:
            print("Flower B")

The two problems I try to solve:

  1. Is there a way to create a variable based on the result of multiple if statement?
  2. Is there a way to do it for all elements of dataframe at the same time instead of using iloc?

As a result I plan to have a dataframe with four columns like this:

expected_result = {
  "blue": [1, 1, 0, 0, 1],
  "red": [1, 1, 0, 1, 1],
    "green": [0,0, 2, 0, 1],
    "flower type": ["Flower A", "Flower A", "Flower B", "Flower A", "Flower B"]
}

#load into df:
expected_result = pd.DataFrame(expected_result)
print(expected_result) 

Thank you very much!


Solution

  • Use np.where:

    In [1367]: import numpy as np
    In [1368]: example['flower type'] = np.where(((example.blue.eq(1) | example.red.eq(1)) & example.green.eq(0)), 'Flower A', 'Flower B')
    
    In [1369]: example
    Out[1369]: 
       blue  red  green flower type
    0     1    1      0    Flower A
    1     1    1      0    Flower A
    2     0    0      2    Flower B
    3     0    1      0    Flower A
    4     1    1      1    Flower B