how to make list of lists from pandas dataframe, skipping nan values

user_7620016 0
Viewed 8933

I have a pandas dataframe that looks roughly like

    foo   foo2   foo3  foo4
a   NY    WA     AZ    NaN
b   DC    NaN    NaN   NaN
c   MA    CA     NaN   NaN

I'd like to make a nested list of the observations of this dataframe, but omit the NaN values, so I have something like [['NY','WA','AZ'],['DC'],['MA',CA'].

There is a pattern in this dataframe, if that makes a difference, such that if fooX is empty, the subsequent column fooY will also be empty.

I originally had something like this code below. I'm sure there's a nicer way to do this

A = [[i] for i in subset_label['label'].tolist()]
B = [i for i in subset_label['label2'].tolist()]
C = [i for i in subset_label['label3'].tolist()]
D = [i for i in subset_label['label4'].tolist()]
out_list = []
for index, row in subset_label.iterrows():
out_list.append([row.label, row.label2, row.label3, row.label4])
out_list
pandas dataframe nan list nested
3 Answers

Here's a vectorized version!

original = pd.DataFrame(data={
    'foo': ['NY', 'DC', 'MA'],
    'foo2': ['WA', np.nan, 'CA'],
    'foo3': ['AZ', np.nan, np.nan],
    'foo4': [np.nan] * 3,
})

out = original.copy().fillna('NAN')

# Build up mapping such that each non-nan entry is mapped to [entry]
#   and nan entries are mapped to []
unique_entries = np.unique(out.values)
mapping = {e: [e] for e in unique_entries}
mapping['NAN'] = []

# Apply mapping
for c in original.columns:
    out[c] = out[c].map(mapping)

# Concatenate the lists along axis 1
out.sum(axis=1)

You should get something like

0    [NY, WA, AZ]
1            [DC]
2        [MA, CA]
dtype: object
user_21086900
Related