Pandas create columns and add values

Viewed 38

I have a pandas Dataframe pmap:

    name            quality_of_life_index
0   Afghanistan     NaN
1   Aland Islands   NaN
2   Albania         98.710473
3   Alderney        NaN
4   Algeria         98.686111

And a function get_country(country) that returns short country name and continent:

get_country("Albania")

>> ('AL', 'EU')

I want to add two columns to pmap and fill them with these values:

    name            quality_of_life_index  country    continent
0   Afghanistan     NaN                    Unknown    Unknown
1   Aland Islands   NaN                    Unknown    Unknown
2   Albania         98.710473              AL         EU
3   Alderney        NaN                    Unknown    Unknown
4   Algeria         98.686111              DZ         AF

I tried this

 for i in range(0, len(pmap)):
     names = get_continent(pmap['name'][i])
     pmap['country', 'continent'] = names

But it returns an error

ValueError: Length of values (2) does not match length of index (233)
2 Answers

You can assign both list with double [[]] with list comprehension:

 pmap[['country', 'continent']] = [get_continent(x) for x in pmap['name']]

You could apply a function for each row as follows:

import pandas as pd
import numpy as np

def get_country(country):
    if "Albania" == country:
        return ('AL', 'EU')
    elif "Algeria" == country:
        return ("DZ", "AF")
    return ("Unknown", "Unknown")

if __name__ == '__main__':
    input = {"name" : ["Afghanistan", "Aland Islands", "Albania", "Alderney", "Algeria"],
    "quality_of_life_index": [np.nan, np.nan, 98.710473, np.nan, 98.686111]}
    df = pd.DataFrame(input)
    df[["country", "continent"]] = df.apply(lambda row: get_country(row["name"]), axis=1, result_type="expand")
    print(df)

Result:

            name  quality_of_life_index  country continent
0    Afghanistan                    NaN  Unknown   Unknown
1  Aland Islands                    NaN  Unknown   Unknown
2        Albania              98.710473       AL        EU
3       Alderney                    NaN  Unknown   Unknown
4        Algeria              98.686111       DZ        AF
Related