JavaScript
x
10
10
1
import pandas as pd
2
3
data = {"Country": ["AA", "BB","CC","DD","EE","FF","GG"],
4
"1990": [0,1,1,1,0,1,1],
5
"1991": [0,0,1,1,1,0,1],
6
"1992": [1,1,1,1,1,0,0],
7
"1993": [0,1,1,1,1,0,0]}
8
9
df = pd.DataFrame(data)
10
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:
JavaScript
1
3
1
for i in range(1,5):
2
print(df[(df == 1)].iloc[:7,0].to_list())
3
I got the output as 4 lists of nans… The desired output would be
JavaScript
1
5
1
c1990=["BB", "CC", "DD", "FF", "GG"]
2
c1991=["CC", "DD", "EE", "GG"]
3
c1992=["AA", "BB","CC","DD","EE"]
4
c1993=["BB","CC","DD","EE"]
5
Advertisement
Answer
One way using dict comprehension with groupby
on axis=1
:
JavaScript
1
9
1
res = {name: i.index[i[name]].tolist() for name, i in df.set_index("Country").astype(bool).groupby(level=0, axis=1)}
2
3
print (res)
4
5
{'1990': ['BB', 'CC', 'DD', 'FF', 'GG'],
6
'1991': ['CC', 'DD', 'EE', 'GG'],
7
'1992': ['AA', 'BB', 'CC', 'DD', 'EE'],
8
'1993': ['BB', 'CC', 'DD', 'EE']}
9