Dataframe with intersection among rows in pandas dataframe?

Viewed 111

I have a challenge in a pandas dataframe. Basically, I have 2 columns. In the first one, I have 3 different classes and in the second a list of students that are enrolled in the subject. The example is as follow:

df = pd.DataFrame({'Class': ['1A', '2B', '2C'],
                    'Students': [['Alice', 'Philips', 'John'],
                                 ['Philips', 'John', 'Anna', 'William'],
                               ['Arthur', 'Alice', 'Anna', 'William']]
                  })

I would like to have a second dataframe with the number of students that are presented in more thant one class. In other words, the intersection between the classes, as follow

result= pd.DataFrame({'Comparison': ['1A-2B','1A-2C', '2B-2C'],
                      'Intersection size': [2, 1, 2]})

Thank you for your help and attention!

2 Answers

You can try the following:

import pandas as pd

df = pd.DataFrame({
    'Class': ['1A', '2B', '2C'],
    'Students': [['Alice', 'Philips', 'John'],
                 ['Philips', 'John', 'Anna', 'William'],
                 ['Arthur', 'Alice', 'Anna', 'William']]
})

ddf = df.explode("Students")
ddf = pd.crosstab(ddf["Students"], ddf["Class"])
A = ddf.values
result = pd.DataFrame(A.T @ A, index=ddf.columns, columns=ddf.columns)
print(result)

It gives:

Class  1A  2B  2C
Class            
1A      3   2   1
2B      2   4   2
2C      1   2   4

The intersection of every row and column gives the number of students taking both classes. Diagonal entries give numbers of students in each individual class.

If you want to get a dataframe listing only combinations of different classes with non-zero intersection values, then the following should work:

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'Class': ['1A', '2B', '2C'],
    'Students': [['Alice', 'Philips', 'John'],
                 ['Philips', 'John', 'Anna', 'William'],
                 ['Arthur', 'Alice', 'Anna', 'William']]
})

ddf = df.explode("Students")
ddf = pd.crosstab(ddf["Students"], ddf["Class"])
A = ddf.values
result = pd.DataFrame(np.tril(A.T @ A, k=-1).T,
                      index=ddf.columns,
                      columns=ddf.columns).stack()
result.index = result.index.map(lambda x: f"{x[0]}-{x[1]}")
result[result > 0]

It gives:

1A-2B    2.0
1A-2C    1.0
2B-2C    2.0
import pandas as pd
import itertools

df = pd.DataFrame({'Class': ['1A', '2B', '2C'],
                    'Students': [['Alice', 'Philips', 'John'],
                                 ['Philips', 'John', 'Anna', 'William'],
                               ['Arthur', 'Alice', 'Anna', 'William']]
                  })
  1. combination: to generate a combination of all colunms. use itertools.
col=list(itertools.combinations(df.Class,2))  
> col  Out[68]:
> [('1A','2B'), ('1A', '2C'), ('2B', '2C')]
  1. explode: to form a structured dataframe

df1=df.explode('Students')

  1. write a for
d={}
for c in col:
    tmp=df1[(df1['Class']==c[0]) | (df1['Class']==c[1])]
    count=len(tmp)-tmp.Students.nunique()
    d[str(c[0])+'-'+str(c[1])]=count

The dictionary d has what you want:

d
Out[71]: {'1A-2B': 2, '1A-2C': 1, '2B-2C': 2}
Related