Flatten a column with value of type list while duplicating the other column's value accordingly in Pandas

StackOverflow https://stackoverflow.com/questions/21160134

Domanda

Dear power Pandas experts:

I'm trying to implement a function to flatten a column of a dataframe which has element of type list, I want for each row of the dataframe where the column has element of type list, all columns but the designated column to be flattened will be duplicated, while the designated column will have one of the value in the list.

The following illustrate my requirements:

input = DataFrame({'A': [1, 2], 'B': [['a', 'b'], 'c']})
     A   B
0    1   [a, b]
1    2   c

expected = DataFrame({'A': [1, 1, 2], 'B': ['a', 'b', 'c']}, index=[0, 0, 1])

     A   B
0    1   a
0    1   b
1    2   c

I feel that there might be an elegant solution/concept for it, but I'm struggling.

Here is my attempt, which does not work yet.

def flattenColumn(df, column):
    '''column is a string of the column's name.
    for each value of the column's element (which might be a list), duplicate the rest of columns at the correspdonding row with the (each) value.
    '''
    def duplicate_if_needed(row):
        return concat([concat([row.drop(column, axis = 1), DataFrame({column: each})], axis = 1) for each in row[column][0]])
    return df.groupby(df.index).transform(duplicate_if_needed)

In recognition of alko's help, here is my trivial generalization of the solution to deal with more than 2 columns in a dataframe:

def flattenColumn(input, column):
    '''
    column is a string of the column's name.
    for each value of the column's element (which might be a list),
    duplicate the rest of columns at the corresponding row with the (each) value.
    '''
    column_flat = pandas.DataFrame(
        [
            [i, c_flattened]
            for i, y in input[column].apply(list).iteritems()
            for c_flattened in y
        ],
        columns=['I', column]
    )
    column_flat = column_flat.set_index('I')
    return (
        input.drop(column, 1)
             .merge(column_flat, left_index=True, right_index=True)
    )

The only limitation at the moment is that the order of columns changed, the column flatten would be at the right most, not in its original position. It should be feasible to fix.

È stato utile?

Soluzione

I guess easies way to flatten list of lists would be a pure python code, as this object type is not well suited for pandas or numpy. So you can do it with for example

>>> b_flat = pd.DataFrame([[i, x] 
...               for i, y in input['B'].apply(list).iteritems() 
...                    for x in y], columns=list('IB'))
>>> b_flat = b_flat.set_index('I')

Having B column flattened, you can merge it back:

>>> input[['A']].merge(b_flat, left_index=True, right_index=True)
   A  B
0  1  a
0  1  b
1  2  c

[3 rows x 2 columns]

If you want the index to be recreated, as in your expected result, you can add .reset_index(drop=True) to last command.

Altri suggerimenti

You can use df.explode(). Check out this method here

It is surprising that there isn't a more "native" solution. Putting the answer from @alko into a function is easy enough:

def unnest(df, col, reset_index=False):
    import pandas as pd
    col_flat = pd.DataFrame([[i, x] 
                       for i, y in df[col].apply(list).iteritems() 
                           for x in y], columns=['I', col])
    col_flat = col_flat.set_index('I')
    df = df.drop(col, 1)
    df = df.merge(col_flat, left_index=True, right_index=True)
    if reset_index:
        df = df.reset_index(drop=True)
    return df

Then simply

input = pd.DataFrame({'A': [1, 2], 'B': [['a', 'b'], 'c']})
expected = unnest(input, 'B')

I guess it would be nice to allow unnesting of multiple columns at once and to handle the possibility of a nested column named I, which would break this code.

A slightly simpler / more readable solution than the ones above which worked for me.

 out = []
 for n, row in df.iterrows():
    for item in row['B']:
        row['flat_B'] = item
        out += [row.copy()]


flattened_df = pd.DataFrame(out)

How about

input = pd.DataFrame({'A': [1, 2], 'B': [['a', 'b'], 'c']})

input[['A', 'B']].set_index(['A'])['B'].apply(pd.Series).stack().reset_index(level=1, drop=True).reset_index().rename(columns={0:'B'})
Out[1]: 
   A  B
0  1  a
1  1  b
2  2  c

One liner - applying the pd.DataFrame constructor, concatenating and joining to original.

my_df = pd.DataFrame({'a': [1, 2, 3], 'b': [2, 3, 4], 'c': [(1, 2), (1, 2), (2, 3)]})
my_df.join(pd.concat(map(lambda x: pd.DataFrame(list(x)), my_df['c']), axis=0))

You can also manipulate the list first, then create a new dataframe: for example:

input = DataFrame({'A': [1, 2], 'B': [['a', 'b'], 'c']})
listA=input.A.tolist()
listB=input.B.tolist()
count_sublist_len=[len(ele) for ele in listB if type(ele)==list else 1]
# create similar list for A
new_listA=[count_sublist_len[i]*[listA[i]] for i in range(len(listA)]
# flatten them
f_A=[item for sublist in new_listA for item in sublist]
f_B=[item for sublist in listB for item in sublist]
df_new=pd.DataFrame({'A':f_A,'B':f_b})

Basically the same as what yaiir did but then using list comprehension in a nice function:

def flatten_col(df: pd.DataFrame, col_from: str, col_to: str) -> pd.DataFrame:
    return pd.DataFrame([row.copy().set_value(col_to, x)
                         for i, row in df.iterrows()
                         for x in row[col_from]]) \
        .reset_index(drop=True)

where col_from is the column containing the lists and col_to is the name of the new column with the split list values.

Use as flatten_col(input, 'B', 'B') in your example. The benefit of this method is that copies along all other columns as well (unlike some other solutions). However it does use the deprecated set_value method..

Autorizzato sotto: CC-BY-SA insieme a attribuzione
Non affiliato a StackOverflow
scroll top