Combination of rows in numpy.ndarray

Tags: , , , ,



I have the following numpy.ndarray

S=np.array([[[ -0.6,  -0.2,   0. ],
             [-60. ,   2. ,   0. ],
             [  6. , -20. ,   0. ]],

            [[ -0.4,  -0.8,   0. ],
             [-40. ,   8. ,   0. ],
             [  4. , -80. ,   0. ]]])

I want to find all the possible combinations of sum of each row (sum of individual elements of a row except the last column) of S[0,:,:] with each row of S[1,:,:], i.e., my desired result is (order does not matter):

array([[-1,     -1],
       [-40.6,   7.8],
       [3.4,   -80.2],
       [-60.4,   1.2],
       [-100,    10],
       [-56,    -78],
       [5.6,   -20.8],
       [-34,    -12],
       [10,    -100]])

which is a 9-by-2 array resulting from 9 possible combinations of S[0,:,:] and S[1,:,:]. Although I have used a particular shape of S here, the shape may vary, i.e., for

x,y,z = np.shape(S)

in the above problem, x=2, y=3, and z=3, but these values may vary. Therefore, I am seeking for a generalized version.

Your help will be highly appreciated. Thank you for your time!

(Please no for loops if possible. It is pretty trivial then.)

Answer

You can use broadcast like this:

(S[0,:,None, :-1] +  S[1,None,:,:-1]).reshape(-1,2)

Output:

array([[  -1. ,   -1. ],
       [ -40.6,    7.8],
       [   3.4,  -80.2],
       [ -60.4,    1.2],
       [-100. ,   10. ],
       [ -56. ,  -78. ],
       [   5.6,  -20.8],
       [ -34. ,  -12. ],
       [  10. , -100. ]])


Source: stackoverflow