In Pandas, count rows that match multiple criteria based on values within the same DataFrame, and add the count in column

user_2859206 0
Viewed 576

I have a DataFrame with an index, and then a reference to other indexes, and an organization. For example:

import pandas as pd
inp = [{'index' :1, 'refIndex':3, 'org' : 'org1'}, {'index':2, 'refIndex':1, 'org': 'org1'}, {'index':3, 'refIndex': 2, 'org' : 'org2'}]
df = pd.DataFrame(inp)
print df

Output:

   index   refIndex   org
0  1       3          org1
1  2       1          org1
2  3       2          org2

What I need to do is count on each row how many other rows where that row's index occurs as the refIndex from the same org.

So I end up with a DataFrame like:

   index   refIndex   org    count
0  1       3          org1   1        # index 1 org1 occurs as refIndex and org once elsewhere
1  2       1          org1   0        # index 2 org1 occurs as refIndex and org nowhere else
2  3       2          org2   0        # index 3 org2 occurs as refIndex and org nowhere else

I am new to Python and Pandas, so please excuse if this is obvious to you. I have been struggling all day with trying groupbys, functions, for loops inside for loops, merges.

pandas dataframe python
2 Answers

Numpy broadcasting

i, r, o = df.values.T
df['count'] = np.sum((i[:, None] == r) & (o[:, None] == o), axis=1)

Explanations

Create a boolean mask by comparing every value in index column with every value in refIndex column.

>>> (i[:, None] == r)

array([[False,  True, False, False],
       [False, False,  True, False],
       [ True, False, False,  True],
       [False, False, False, False]])

Create another boolean mask by comparing every value in org column with every other value in the org column itself.

>>> (o[:, None] == o)

array([[ True,  True, False, False],
       [ True,  True, False, False],
       [False, False,  True,  True],
       [False, False,  True,  True]])

Take the logical and of the above two boolean masks so that we could be able to mask the True values in first boolean mask which does not belong to the same organisation.

>>> (i[:, None] == r) & (o[:, None] == o)

array([[False,  True, False, False],
       [False, False, False, False],
       [False, False, False,  True],
       [False, False, False, False]])

Finally take sum along axis=1 to count of rows where the row's index occurs as the refIndex from the same org.

>>> df

   index  refIndex   org  count
0      1         3  org1      1
1      2         1  org1      0
2      3         2  org2      0
user_128331660

Try

df = df.rename(columns={'index': 'id'})
groups = df.groupby('org')
df.apply(lambda x: x.refIndex not in groups.get_group(x.org)['id'], axis=1).astype(int)
user_51200490
Related