Pandas groupby and count unique value of column

Viewed 280

I'm having some difficulties with a Pandas query and would appreciate some help. The DataFrame I'm using is the following:

data = {
    "id": ["A", "A", "A", "A", "A", "A", "B", "B"],
    "start": ["P1", "P2", "P3", "P4", "P1", "P5", "P1", "P3"],
    "finish": ["P2", "P3", "P4", "P1", "P5", "P1", "P3", "P5"],
    "order": [1, 2, 3, 4, 5, 6, 1, 2]
}
df = pd.DataFrame(data)
print(df)

  id start finish  order
0  A    P1     P2      1
1  A    P2     P3      2
2  A    P3     P4      3
3  A    P4     P1      4
4  A    P1     P5      5
5  A    P5     P1      6
6  B    P1     P3      1
7  B    P3     P5      2

For each id, I would like to calculate the total number of unique points reached in the finish column, excluding the paths that start from P1.

If we take a closer look at id A, the expected result should be 3:

Step1: P1 -> P2 [X] (starting in P1 is forbidden) 
Step2: P2 -> P3 +1  
Step3: P3 -> P4 +1  
Step4: P4 -> P1 +1  
Step5: P1 -> P5 [X] (starting in P1 is forbidden)
Step6: P5 -> P1 +0  (P1 already visited in step 4)

Using the same logic for id B, the expected result should be 1. Thus, the resulting DataFrame I'm looking is the following:

  id result
0  A      3 
1  B      1    

Here are my thoughts so far: I started by creating a mask to filter out all the records where the starting point is P1. Then, I need to group the remaining records by id and count the unique values in the finish column. My issue is that I don't know how to transcribe what I just explained in Python/Pandas language. Any help would be appreciated.

1 Answers

We can drop all lines with start=='P1', then groupby id and count unique finish:

(df[df['start'].ne('P1')]       # drop rows with `start` == 'P1'
   .groupby('id')               # group by `id`
   ['finish'].nunique()         # count unique `finish`
   .reset_index(name='result')  # match the output
)

Output:

  id  result
0  A       3
1  B       1
Related