Split a pipe-delimited series, groupby a separate series, and return the counts of each split value in new columns

user_9245853 0
Viewed 402

Given a dataframe with a pipe-delimited series:

import pandas as pd
import matplotlib.pyplot as plt

df = pd.DataFrame({'year': [1960, 1960, 1961, 1961, 1961],
                   'genre': ['Drama|Romance|Thriller',
                             'Spy|Mystery|Bio',
                             'Drama|Romance',
                             'Drama|Romance',
                             'Drama|Spy']})

Or in data format:

   year                   genre
0  1960  Drama|Romance|Thriller
1  1960         Spy|Mystery|Bio
2  1961           Drama|Romance
3  1961           Drama|Romance
4  1961               Drama|Spy

I can split the genre series with str.split (as demonstrated in many similar questions on SO).

But I would also like to group by the year and return the count of Drama, Romance, Thriller, and so on for each unique year in new columns.

My initial attempt:

df_split = df.groupby('year')['genre'].apply(lambda x: x.str.split('|', expand=True).reset_index(drop=True))

which returns

            0        1         2
year                            
1960 0  Drama  Romance  Thriller
     1    Spy  Mystery       Bio
1961 0  Drama  Romance       NaN
     1  Drama  Romance       NaN
     2  Drama      Spy       NaN

but then how to get the count of each unique genre in its own column, by year?

I can get the unique genres using

genres = pd.unique(df['genre'].str.split('|', expand=True).stack())

but am still unsure of how to get the genres as separate series, with their counts by year.

The final output I'd like is:

      Drama  Romance  Thriller  Spy  Mystery  Bio
1960      1        1         1    1        1    1
1961      3        2         0    1        0    0

where each unique genre is its own series, with its corresponding count by year.

This may also very well also be an X-Y problem. My end goal is to produce a percentage stacked-area chart. Assuming df_split has the required transformation, I'd like to do:

df_perc = df_split.divide(df_split.sum(axis=1), axis=0)

which returns

         Drama   Romance  Thriller       Spy   Mystery       Bio
1960  0.166667  0.166667  0.166667  0.166667  0.166667  0.166667
1961  0.500000  0.333333  0.000000  0.166667  0.000000  0.000000

and then

plt.stackplot(df_perc.index, *[ts for col, ts in df_perc.iteritems()],
                               labels=df_perc.columns)
plt.gca().set_xticks(df_perc.index)
plt.margins(0)
plt.legend()

giving the output:

enter image description here

pandas matplotlib python
2 Answers

We can get to your desired result using some simple reshaping and aggregation:

(df.assign(genre=df['genre'].str.split('|'))
   .explode('genre')
   .groupby('year')['genre']
   .value_counts(normalize=True)
   .unstack(fill_value=0))     
 
genre       Bio     Drama   Mystery   Romance       Spy  Thriller
year                                                             
1960   0.166667  0.166667  0.166667  0.166667  0.166667  0.166667
1961   0.000000  0.500000  0.000000  0.333333  0.166667  0.000000

From here you can finish up by plotting an area plot:

(df.assign(genre=df['genre'].str.split('|'))
   .explode('genre')
   .groupby('year')['genre']
   .value_counts(normalize=True)
   .unstack(fill_value=0)
   .plot
   .area())

How It Works

Start by exploding your data across rows:

df.assign(genre=df['genre'].str.split('|')).explode('genre') 

   year     genre
0  1960     Drama
0  1960   Romance
0  1960  Thriller
1  1960       Spy
1  1960   Mystery
1  1960       Bio
2  1961     Drama
2  1961   Romance
3  1961     Drama
3  1961   Romance
4  1961     Drama
4  1961       Spy

Next, do a groupby and get the normalized count:

_.groupby('year')['genre'].value_counts(normalize=True)

year  genre   
1960  Bio         0.166667
      Drama       0.166667
      Mystery     0.166667
      Romance     0.166667
      Spy         0.166667
      Thriller    0.166667
1961  Drama       0.500000
      Romance     0.333333
      Spy         0.166667
Name: genre, dtype: float64

Next, unstack the result:

_.unstack(fill_value=0)

genre       Bio     Drama   Mystery   Romance       Spy  Thriller
year                                                             
1960   0.166667  0.166667  0.166667  0.166667  0.166667  0.166667
1961   0.000000  0.500000  0.000000  0.333333  0.166667  0.000000

Finally, plot with

_.plot.area()
user_49090870

You could re-arrange your data in the first place:

import pandas as pd
from itertools import groupby
from collections import defaultdict

data = """
1960  Drama|Romance|Thriller
1960         Spy|Mystery|Bio
1961           Drama|Romance
1961           Drama|Romance
1961               Drama|Spy
"""

# sort it first by year
lst = sorted((line.split() for line in data.split("\n") if line), key=lambda x: x[0])

# group it by year, expand the genres
result = {}
for key, values in groupby(lst, key=lambda x: x[0]):
    dct = defaultdict(int)
    for lst in values:
        for genre in lst[1].split("|"):
            dct[genre] += 1
    result[key] = dct

# feed it all to pandas
df = pd.DataFrame.from_dict(result, orient='index').fillna(0)

print(df)

Which would yield

      Drama  Romance  Thriller  Spy  Mystery  Bio
1960      1        1       1.0    1      1.0  1.0
1961      3        2       0.0    1      0.0  0.0
user_12314500
Related