Search code examples
pythonnumpycombinationspython-itertoolsadjacency-matrix

Generate adjacency matrix from a list, where adjacency means equal elements


I have a list like this:

lst = [0, 1, 0, 5, 0, 1]

I want to generate an adjacency matrix:

out = 
array([[ 1.,  0.,  1.,  0.,  1.,  0.],
       [ 0.,  1.,  0.,  0.,  0.,  1.],
       [ 1.,  0.,  1.,  0.,  1.,  0.],
       [ 0.,  0.,  0.,  1.,  0.,  0.],
       [ 1.,  0.,  1.,  0.,  1.,  0.],
       [ 0.,  1.,  0.,  0.,  0.,  1.]])

where out[i,j] = 1 if lst[i]==lst[j]

Here is my code with two for loops:

lst = np.array(lst)
label_lst = list(set(lst))
out = np.eye(lst.size, dtype=np.float32)
for label in label_lst:
  idx = np.where(lst == label)[0]
  for pair in itertools.combinations(idx,2):
    out[pair[0],pair[1]] = 1
    out[pair[1],pair[0]] = 1

But I feel there should be a way to improve this. Any suggestion?


Solution

  • Use broadcasted comparison -

    np.equal.outer(lst, lst).astype(int) # or convert to float
    

    Sample run -

    In [787]: lst = [0, 1, 0, 5, 0, 1]
    
    In [788]: np.equal.outer(lst, lst).astype(int)
    Out[788]: 
    array([[1, 0, 1, 0, 1, 0],
           [0, 1, 0, 0, 0, 1],
           [1, 0, 1, 0, 1, 0],
           [0, 0, 0, 1, 0, 0],
           [1, 0, 1, 0, 1, 0],
           [0, 1, 0, 0, 0, 1]])
    

    Or convert to array and then manually extend to 2D and compare -

    In [793]: a = np.asarray(lst)
    
    In [794]: (a[:,None]==a).astype(int)
    Out[794]: 
    array([[1, 0, 1, 0, 1, 0],
           [0, 1, 0, 0, 0, 1],
           [1, 0, 1, 0, 1, 0],
           [0, 0, 0, 1, 0, 0],
           [1, 0, 1, 0, 1, 0],
           [0, 1, 0, 0, 0, 1]])