How to efficiently select several value ranges in Pandas?

Viewed 246

I was revising some old code and stumbled upon this snippet.

df = pd.DataFrame({
'x': ['a', 'b', 'c', 'd', 'e', 'dc', 'ca', 'cd', 'cf', 'cv', 'cs', 'ca', 'ac', 'fc'],
'a': [34, 28, 51,5,120,12,45,56,67,54,34,32,1213,2]})

five = df[df['a'] < 5]
ten = df[(df['a'] > 5)  &  (df['a'] < 10)]
twenty = df[(df['a'] > 10)  &  (df['a'] < 20 )]
thirty = df[(df['a'] > 20)  &  (df['a'] < 30 )]
forty = df[(df['a'] > 30)  &  (df['a'] < 40 )]
fifty = df[(df['a'] > 40)  &  (df['a'] < 50 )]
sixty =  df[(df['a'] > 50)  &  (df['a'] < 60)]
over =  df[(df['a'] > 60)]

Basically there is a DataFrame and I need to group values that are within a certain range. There are multiple ranges. The grouped values are later used in a box plot. The code above gets the job done but I am pretty sure there is a better way to do this!

Question: What if I need to create 1000 groups? I want to change the edge values and add new groups more dynamically. How can this be done?

3 Answers

You can use pd.cut() to create bins and then df.groupby() to create a grouper object with list of dataframes split by bin. Then you can store that in a list.

(Note: I have not named each of the dfs because as you say if you have 1000 dfs, it doesn't make sense to name each. You can just store them in a list.)

import pandas as pd

cuts = [5,10,20,30,40,50,60]   #pd.cut take 2 elements from this for grouping (5,10); (30,40) etc
cols = ['x','a']             #columns to keep in final list of dataframes

df1['bin'] = pd.cut(df1['a'], cuts, right=False)   #create bins
group = df1.groupby('bin')   #groupby bins (grouper object)
dfs = [j[cols] for i,j in group]   #Iterate over grouper object to get dfs


#Printing block (ignore this)
k = list(zip(cuts,cuts[1:]))
for i,j in enumerate(dfs):
    print('Group in range',k[i],':')
    print(j)
    print('\n')
Group in range (5, 10) :
   x  a
3  d  5


Group in range (10, 20) :
    x   a
5  dc  12


Group in range (20, 30) :
   x   a
1  b  28


Group in range (30, 40) :
     x   a
0    a  34
10  cs  34
11  ca  32


Group in range (40, 50) :
    x   a
6  ca  45


Group in range (50, 60) :
    x   a
2   c  51
7  cd  56
9  cv  54    

You can check the documentation for cut and groupby for more details

METHOD#1

You can use pd.cut and the create dynamic groups and save them in a dictionary, ,the refer each keys for the individual dataframe:

bins = [0,5,10,20,30,40,50,60,np.inf]
labels = ['five','ten','twenty','thirty','forty','fifty','sixty','over']

u = df1.assign(grp=pd.cut(df1['a'],bins,labels=labels))
d = dict(iter(u.groupby("grp")))

test runs:

print(f"""Group five is \n\n {d['five']}\n\n 
         Group forty is \n\n{d['forty']} \n\n Group over is \n\n{d['over']}""")

Group five is 

      x  a   grp
3    d  5  five
13  fc  2  five

 
Group forty is 

     x   a    grp
0    a  34  forty
10  cs  34  forty
11  ca  32  forty 

 Group forty is 

     x     a   grp
4    e   120  over
8   cf    67  over
12  ac  1213  over

METHOD#2 you can also use locals for saving dictionary keys a local variables but the dict method is better:

bins = [0,5,10,20,30,40,50,60,np.inf]
labels = ['five','ten','twenty','thirty','forty','fifty','sixty','over']

u = df1.assign(grp=pd.cut(df1['a'],bins,labels=labels))
d = dict(iter(u.groupby("grp")))
for k,v in d.items():
    locals().update({k:v})

print(over,'\n\n',five,'\n\n',sixty)

     x     a   grp
4    e   120  over
8   cf    67  over
12  ac  1213  over 

      x  a   grp
3    d  5  five
13  fc  2  five 

     x   a    grp
2   c  51  sixty
7  cd  56  sixty
9  cv  54  sixty

You could use pandas.cut() to assign certain bins to your values. For example: you can add an extra column category in your dataframe as such:

df['category'] = pd.cut(df['a'], bins=np.arange(5, 65, 5))

This will give you the following dataframe:

     x     a      category
0    a    34  (30.0, 35.0]
1    b    28  (25.0, 30.0]
2    c    51  (50.0, 55.0]
3    d     5           NaN
4    e   120           NaN
5   dc    12  (10.0, 15.0]
6   ca    45  (40.0, 45.0]
7   cd    56  (55.0, 60.0]
8   cf    67           NaN
9   cv    54  (50.0, 55.0]
10  cs    34  (30.0, 35.0]
11  ca    32  (30.0, 35.0]
12  ac  1213           NaN
13  fc     2           NaN

and then use those values in a groupby() method using the new column after that. The bins that you supply as argument can be created quite flexibly (as shown above).

The documentation for the cut() function is here: https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.cut.html

Related