Given a matrix, return a list (or generator) of all it's pairwise-neighbors-CodePudding

Let's say i have this matrix: (doesn't have to be squared)

a b c
d e f
g h i

I want to return a list, or a generator, of all it's pairwise-neighbors. Meaning:

[[a,b], [b,c], [d,e], [e,f], [g,h], [h,i], [a,d], [b,e], [c,f], [d,g], [e,h], [f,i]]

AND(!) to add an option to return only the sum of: a*b b*c d*e e*f g*h ... f*i

If this is not explained clearly enough from the example, two elements are pairwise-neighbors if they both stand next to each other, from left to right or from the bottom up (NOT diagonally!).

Here's what I did so far:

def neighbors(X, returnSet=False):
    '''
    :param X: The matrix
    :param returnSet:  False[DEFAULT] = Return sum of Xi*Xj (when Xi and Xj are pairwise-neighbors)
                       True = Return the set of all neighbors.
    :return: Depends on returnSet
    '''
    sum, neighbors = 0, []
    for i in range(0, X.shape[0]):
        for j in range(0, X.shape[1]):
            if j   1 < X.shape[1]:
                neighbors.append([X[i][j], X[i][j   1]]) if returnSet else None
                sum  = X[i][j] * X[i][j   1]
            if i   1 < X.shape[0]:
                neighbors.append([X[i][j], X[i   1][j]]) if returnSet else None
                sum  = X[i][j] * X[i   1][j]
    return neighbors if returnSet else sum

This code works, but I don't really like how it looks. Can you come up with some "cooler" algorithm than that? Or maybe more efficient? I like when a python code are short and simple.

CodePudding user response：

Basically same as your last question, just doing it over two axes and with a double dot product (which isn't as pretty).

a = np.arange(9).reshape(3,3)
np.einsum('ij, ij->', a[1:], a[:-1])   np.einsum('ij, ij->', a[:, 1:], a[:, :-1])
Out[]: 232

neighbors(a)
Out[]: 232

If you just want the pairs:

h = np.stack([a[:, 1:], a[:, :-1]]).transpose(1,2,0).reshape(-1, 2)
v = np.stack([a[1:], a[:-1]]).transpose(1,2,0).reshape(-1, 2)
np.vstack([h, v])
Out[]: 
array([[1, 0],
       [2, 1],
       [4, 3],
       [5, 4],
       [7, 6],
       [8, 7],
       [3, 0],
       [4, 1],
       [5, 2],
       [6, 3],
       [7, 4],
       [8, 5]])