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