Cartesian product of grouped arrays in list

I have a data like this,

mydata = [np.array([8,6]),
         np.array([9]),
         np.array([6]),
         np.array([5]),
         np.array([10,3])]

And an array representing the group of the data above.

path = np.array([1,2,1,2,2])
# It can be read as index 0,2 is in group1 while index 1,3,4 is in group2 for mydata.

I want to get the cartesian product with resprect to the groups at the right format. Here is my trial with meshgrid method of the numpy .

import numpy as np
from itertools import compress

out = []
for i in np.unique(path):
    out.append(list(compress(mydata, np.isin(path,i))))



[*np.array(np.meshgrid(*out)).T.reshape(-1,len(out))]

# [array([array([8, 6]), array([9])], dtype=object),
# array([array([8, 6]), array([5])], dtype=object),
#  array([array([8, 6]), array([10,  3])], dtype=object),
#  array([array([6]), array([9])], dtype=object),
#  array([array([6]), array([5])], dtype=object),
#  array([array([6]), array([10,  3])], dtype=object)]

It gives array in array together with dtype. However my expected output is the simpler version of the above:

desired = [np.array([8,6,9]),
           np.array([8,6,5]), 
           np.array([8,6,10,3]),
           np.array([6,9]),
           np.array([6,5]),
           np.array([6,10,3])]

By the way, I am also open to less complicated solutions. Honestly, I did not like my approach.

Thanks in advance.

CodePudding user response：

This is more of a list task than an array one.

In [120]: mydata = [[8, 6], [9], [6], [5], [10, 3]]
     ...: path = [1, 2, 1, 2, 2]
In [121]: out = []
     ...: for i in np.unique(path):
     ...:     out.append(list(compress(mydata, np.isin(path, i))))
     ...: 
In [122]: out
Out[122]: [[[8, 6], [6]], [[9], [5], [10, 3]]]

We can easily construct the cartesian product with a list comprehension. This uses list to join the elements:

In [123]: [x   y for x in out[0] for y in out[1]]
Out[123]: [[8, 6, 9], [8, 6, 5], [8, 6, 10, 3], [6, 9], [6, 5], [6, 10, 3]]

starting with arrays:

In [126]: mydata = [
     ...:     np.array([8, 6]),
     ...:     np.array([9]),
     ...:     np.array([6]),
     ...:     np.array([5]),
     ...:     np.array([10, 3]),
     ...: ]
In [127]: out = []
     ...: for i in np.unique(path):
     ...:     out.append(list(compress(mydata, np.isin(path, i))))
     ...: 
In [128]: out
Out[128]: [[array([8, 6]), array([6])], [array([9]), array([5]), array([10,  3])]]

use hstack to join the arrys:

In [130]: [np.hstack((x,y)) for x in out[0] for y in out[1]]
Out[130]: 
[array([8, 6, 9]),
 array([8, 6, 5]),
 array([ 8,  6, 10,  3]),
 array([6, 9]),
 array([6, 5]),
 array([ 6, 10,  3])]