Pandas fill column by count of other columns

Viewed 486

Example df:

     company   vehicle registration
0   company1     truck       abc123
1   company1     truck      abcdefg
2   company1       car       234cse
3   company1  forklift          NaN
4   company1     truck        93ds2
5   company2       car      rentall
6   company2       car      rental2
7   company2     truck      rentals
8   company2     truck      rental*
9   company2       car      rental5
10  company3     truck       fdsa23
11  company3     truck        asdf4
12  company3     other       fdsag3
13  company3     other          NaN
14  company3     truck      gls319d

sample_data

My goal is to get counts by company and vehicle type (registration and vehicle columns will be dropped).

I've tried this:

import pandas as pd

df = pd.read_csv('path to csv', header=0)

df.loc[df.vehicle == 'truck', 'trucks'] = 1
df.loc[df.vehicle == 'car', 'cars'] = 1
df.loc[df.vehicle != 'truck', 'others'] = 1
df.loc[df.vehicle != 'cars', 'others'] = 1

from there I assume some sort of groupby and sum function would consolidate the rows and columns.

Unfortunately, this only fills in the vehicle columns with a "1" value rather than having the values in the respective columns.

My desired output is:


company   trucks  cars  others
company1  3       1     1 
company2  2       3     0
company3  3       0     2

I'm sure this has probably been answered before, but my google-fu is weak this morning.

Cheers.

1 Answers

First use Series.map by filtered categories in dictionary and replace all no matched values (NaNs) by Series.fillna.

Then pass to crosstab and if order of output columns is important add DataFrame.reindex:

df['new'] = df.vehicle.map({'truck':'trucks', 'car':'cars'}).fillna('other')
df = pd.crosstab(df['company'], df['new']).reindex(['cars','trucks','other'], axis=1)
print (df)
vehicle   cars  trucks  other
company                      
company1     1       3      1
company2     3       2      0
company3     0       3      2
Related