Co-occurrence Matrix from list of words in Python

Question

I have a list of names like:

names = ['A', 'B', 'C', 'D']

and a list of documents, that in each documents some of these names are mentioned.

document =[['A', 'B'], ['C', 'B', 'K'],['A', 'B', 'C', 'D', 'Z']]

I would like to get an output as a matrix of co-occurrences like:

  A  B  C  D
A 0  2  1  1
B 2  0  2  1
C 1  2  0  1
D 1  1  1  0

There is a solution (Creating co-occurrence matrix) for this problem in R, but I couldn't do it in Python. I am thinking of doing it in Pandas, but yet no progress!

Answer 1

Here is another solution using itertools and the Counter class from the collections module.

import numpy
import itertools
from collections import Counter

document =[['A', 'B'], ['C', 'B'],['A', 'B', 'C', 'D']]

# Get all of the unique entries you have
varnames = tuple(sorted(set(itertools.chain(*document))))

# Get a list of all of the combinations you have
expanded = [tuple(itertools.combinations(d, 2)) for d in document]
expanded = itertools.chain(*expanded)

# Sort the combinations so that A,B and B,A are treated the same
expanded = [tuple(sorted(d)) for d in expanded]

# count the combinations
c = Counter(expanded)


# Create the table
table = numpy.zeros((len(varnames),len(varnames)), dtype=int)

for i, v1 in enumerate(varnames):
    for j, v2 in enumerate(varnames[i:]):        
        j = j + i 
        table[i, j] = c[v1, v2]
        table[j, i] = c[v1, v2]

# Display the output
for row in table:
    print(row)

The output (which could be easilty turned into a DataFrame) is:

[0 2 1 1]
[2 0 2 1]
[1 2 0 1]
[1 1 1 0]