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']}