Skip to content
Advertisement

Combination of rows in numpy.ndarray

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.)

Advertisement

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. ]])
User contributions licensed under: CC BY-SA
10 People found this is helpful
Advertisement