Skip to content
Advertisement

Filter column value from other columns’ values and turn the results into multiple lists Pandas

   import pandas as pd

   data = {"Country": ["AA", "BB","CC","DD","EE","FF","GG"],
        "1990": [0,1,1,1,0,1,1],
        "1991": [0,0,1,1,1,0,1],
        "1992": [1,1,1,1,1,0,0],
        "1993": [0,1,1,1,1,0,0]}

   df = pd.DataFrame(data)

The goal is: for column 1990-1993, if value == 1, return Country to 4 lists, I also want to set each list a #name of the year and don’t know how to do that.

here is my try:

    for i in range(1,5):
        print(df[(df == 1)].iloc[:7,0].to_list())

I got the output as 4 lists of nans… The desired output would be

c1990=["BB", "CC", "DD", "FF", "GG"]
c1991=["CC", "DD", "EE", "GG"]
c1992=["AA", "BB","CC","DD","EE"]
c1993=["BB","CC","DD","EE"]

Advertisement

Answer

One way using dict comprehension with groupby on axis=1:

res = {name: i.index[i[name]].tolist() for name, i in df.set_index("Country").astype(bool).groupby(level=0, axis=1)}

print (res)

{'1990': ['BB', 'CC', 'DD', 'FF', 'GG'],
 '1991': ['CC', 'DD', 'EE', 'GG'],
 '1992': ['AA', 'BB', 'CC', 'DD', 'EE'],
 '1993': ['BB', 'CC', 'DD', 'EE']}
User contributions licensed under: CC BY-SA
4 People found this is helpful
Advertisement