1

This is almost the same as this question, but in numpy, and for matrices over some dimension:

I have two matrices of length n, A and B, for example,

n = 1000
A = np.random.rand(n, 3, 4, 5)
B = np.random.rand(n, 3, 4, 5)

I want a new matrix C, in which every element (C[i]) is a matrix of shape (2, 3, 4, 5): C[i][0] is A[k], and C[i][0] is B[k] for some k, where elements from A and B are selected once for every other element in the other.

A more concrete example:

A = [elem1, elem2, elem3]
B = [elem4, elem5, elem6]

then

C = [
[elem1, elem4],[elem1, elem5],[elem1, elem6],
[elem2, elem4],[elem2, elem5],[elem2, elem6],
[elem3, elem4],[elem3, elem5],[elem3, elem6]
]

I am aware I can use itertools to build this, but i figured it already exists, and faster.

Gulzar
  • 23,452
  • 27
  • 113
  • 201
  • Check the solutions here: https://stackoverflow.com/questions/58971397/for-loop-using-more-than-one-list-in-python It might help you! – LocoGris Nov 21 '19 at 16:13
  • 2
    What you're describing is the Cartesian product, not the cross product (might be the reason you didn't find the solution yourself). [This beautiful SO post](https://stackoverflow.com/questions/11144513/cartesian-product-of-x-and-y-array-points-into-single-array-of-2d-points) offers a solution with numpy-builtins as well as a bountied implementation of what the author calls an "almost canonical Cartesian product" – Lukas Thaler Nov 21 '19 at 16:15
  • Did the posted solution work for you? – Divakar Nov 22 '19 at 12:52
  • @divakar will check on sunday when im back at work. Dont worry – Gulzar Nov 22 '19 at 13:59
  • @Gulzar Gentle reminder - Did the posted solution work for you? – Divakar Nov 25 '19 at 04:43

1 Answers1

1

From a simpler understanding point of view, an assignment based would make sense. Hence -

dt = np.result_type(A.dtype,B.dtype)
s = A.shape
C = np.empty((s[0],)+(2,)+s[1:],dtype=dt)
C[:,0] = A
C[:,1] = B

This could be simplified with a stacking operation for a one-liner -

C = np.stack((A,B),axis=1)

Bit of explanation on how these two methods solve it :

The one-liner is basically stacking A after B. In the first approach, the resultant is a 2 element ndarray along the second axis - s[0],)+(2,)+s[1:]. Thus, each element off A and off B are interleaved along that axis. This stacking operation has a NumPy builtin np.stack, which is the second method. Hence, these two are equivalent and solve our case.

And some verification :

Question is asking for - a new matrix C, in which every element (C[i]) is a matrix of shape (2, 3, 4, 5). Hence, C would be of shape (n,2,3,4,5). The stacking is along the second axis. Furthermore, from what I have understood, the question is asking for C[k][0] == A[k] for k=0..n and C[k][1] == B[k] for k=0..n. Let's prove that -

In [323]: n = 1000
     ...: A = np.random.rand(n, 3, 4, 5)
     ...: B = np.random.rand(n, 3, 4, 5)

In [324]: C = np.stack((A,B),axis=1)

In [325]: np.all([np.allclose(C[i][0],A[i]) for i in range(1000)])
Out[325]: True

In [326]: np.all([np.allclose(C[i][1],B[i]) for i in range(1000)])
Out[326]: True
Divakar
  • 218,885
  • 19
  • 262
  • 358