Graph Transformer Positional Encoding

Problem

Compute a walk-based positional encoding for every node.

The encoding of a node is the number of walks of length exactly k starting at it (summed over all endpoints). Compute it by repeatedly multiplying the all-ones vector by the adjacency matrix k times. edges is a list of [u, v] pairs. Return the encoding values in node order.

Hint

Walk counts give each node a structural signature independent of node labels.

Test cases

case 1

in: [3,[[0,1],[1,2],[2,0]],2]

expected: [4,4,4]

case 2

in: [4,[[0,1],[1,2],[2,3]],1]

expected: [1,2,2,1]

case 3

in: [3,[[0,1],[1,2],[2,0]],1]

expected: [2,2,2]

case 4

in: [2,[],2]

expected: [0,0]

Starter codePython

def positional_encoding(n, edges, k):
    # Your code here
    pass

4 test cases · run your code against them in the editor