Pairwise combinations of pandas columns

user_6697589 0
Viewed 282

I have a pandas dataframe with a large number of columns, and I would like to generate pairwise counts of columns that meet certain criteria. Given a dataframe like this:

df = pd.DataFrame({'A':['foo','foo','bar','bar'],
'B':['foo','foo','foo','bar'],
'C':['foo','bar','foo','bar'],
'D':['foo','foo','foo','foo'],
'E':['bar','bar','bar','bar']})

I want to create a function that will tell me, for example, the number of rows where two columns are both 'foo' for each combination of columns. I can do this with a looping approach like this:

def pairwise_crit_count(df,crit):
    # define column list
    col_list = df.columns
    # initialize dict that will turn into a pd frame for the result
    result_dict = {'row':col_list}
    for first_col in col_list:
        # create empty list that will become the column of result
        temp_list = []
        for second_col in col_list:
            # count the number of rows that meet the criteria
            num = df.loc[(df[first_col]==crit) & (df[second_col]==crit)].shape[0]
            temp_list.append(num)
        # add to result dict
        result_dict[first_col] = temp_list
    return pd.DataFrame(result_dict)

pairwise_crit_count(df,'foo')

but I feel like there is a cleaner solution that I'm missing.

pandas python
1 Answers

In this case, co-occurences mean matrix multiplication:

s = df.eq('foo').astype(int)
out = s.T @ s

Output:

   A  B  C  D  E
A  2  2  1  2  0
B  2  3  2  3  0
C  1  2  2  2  0
D  2  3  2  4  0
E  0  0  0  0  0
user_42384080
Related