CODEWITHSUNDEEP
pythonpandasdataframe
Hassan Syyid
Hassan Syyid

Reputation: 1581

Pandas create new column with array of other columns with a non null value

I have a dataframe like the following:

   List1  List2  List3  List4  List5  List6  List7  List8
0    NaN    NaN      1    NaN    NaN    NaN      1    NaN
1    NaN    NaN      1    NaN    NaN    NaN      1    NaN
2    NaN    NaN      1    NaN    1.0    NaN      1    NaN
3    NaN    NaN      1    NaN    NaN    NaN      1    NaN
4    NaN    NaN      1    NaN    1.0    NaN      1    NaN

I want to create a new Column called Lists which is an array of all the other columns with a non null value. ie:

                           Lists
0    ['List3', 'List7']
1    ['List3', 'List7']
2    ['List3', 'List5', 'List7']
3    ['List3', 'List7']
4    ['List3', 'List5', 'List7']

I accomplished this with an iterrows() loop, but it's not performant at all. Would appreciate any ideas here.

Upvotes: 0

Views: 578

Answers (3)

Andrej Kesely
Andrej Kesely

Reputation: 195438

Another version:

df["Lists"] = df.apply(lambda x: x[x.notna()].index.tolist(), axis=1)
print(df)

Prints:

   List1  List2  List3  List4  List5  List6  List7  List8                  Lists
0    NaN    NaN      1    NaN    NaN    NaN      1    NaN         [List3, List7]
1    NaN    NaN      1    NaN    NaN    NaN      1    NaN         [List3, List7]
2    NaN    NaN      1    NaN    1.0    NaN      1    NaN  [List3, List5, List7]
3    NaN    NaN      1    NaN    NaN    NaN      1    NaN         [List3, List7]
4    NaN    NaN      1    NaN    1.0    NaN      1    NaN  [List3, List5, List7]

Upvotes: 1

Erfan
Erfan

Reputation: 42916

We can use DataFrame.dot to get all columns which are notna:

df["Lists"] = df.notna().dot(df.columns+",").str.rstrip(",").str.split(",")

   List1  List2  List3  List4  List5  List6  List7  List8                  Lists
0    NaN    NaN      1    NaN    NaN    NaN      1    NaN         [List3, List7]
1    NaN    NaN      1    NaN    NaN    NaN      1    NaN         [List3, List7]
2    NaN    NaN      1    NaN   1.00    NaN      1    NaN  [List3, List5, List7]
3    NaN    NaN      1    NaN    NaN    NaN      1    NaN         [List3, List7]
4    NaN    NaN      1    NaN   1.00    NaN      1    NaN  [List3, List5, List7]

Upvotes: 1

ThePyGuy
ThePyGuy

Reputation: 18416

You can use pandas.DataFrame.apply on axis=1 with a lambda with a list-comprehension and take only the columns that has non-NaN value for each row using pd.notna():

>>> df.apply(lambda x: [c for c in df if pd.notna(x[c])], axis=1)

0           [List3, List7]
1           [List3, List7]
2    [List3, List5, List7]
3           [List3, List7]
4    [List3, List5, List7]
dtype: object

Upvotes: 0

Related Questions