Python List Comprehension with DataFrames – Part 2

Let us continue our tutorial with nested list comprehensions. See the below example:

matrix = [[col for col in range(1,5)] for row in range(1,5)]

for row in matrix:
    print(row)

#Output
#[1, 2, 3, 4]
#[1, 2, 3, 4]
#[1, 2, 3, 4]
#[1, 2, 3, 4]

With the above list comprehension we are able to output a 4×4 array using a nested list comprehension.

We can use if statements or conditionals in our list comprehension also:

import pandas as pd

df = pd.DataFrame({'name': ['John', 'Paul', 'George', 'Ringo'], 
        'Income': [10000, 900, 5566, 100000]})

df['CLASS'] = ["HIGH INCOME" if x > 10000 else "LOW INCOME" for x in df['Income']]

print(df)

#Output
#     name  Income        CLASS
#0    John   10000   LOW INCOME
#1    Paul     900   LOW INCOME
#2  George    5566   LOW INCOME
#3   Ringo  100000  HIGH INCOME