Как в NumPy из k векторов получить (k+1)-мерный массив комбинаций?

Question

[dom1n1k]

Проще всего это объяснить на примере.
Имеется k (в данном случае k=2) векторов одинаковой длины:

x = [1, 2, 3]
y = [10, 20, 30]

Нужно получить матрицу комбинаций, каждый элемент которой будет массивом (то самое +1 измерение):

[
[[1, 10], [1, 20], [1, 30]],
[[2, 10], [2, 20], [2, 30]],
[[3, 10], [3, 20], [3, 30]]
]

Аналогично, если добавить третий вектор
z = [100, 200, 300]
То нужно получить что-то такое:

[[
[[1, 10, 100], [1, 10, 200], [1, 10, 300]],
[[1, 20, 100], [1, 20, 200], [1, 20, 300]],
[[1, 30, 100], [1, 30, 200], [1, 30, 300]]
], [
[[2, 10, 100], [2, 10, 200], [2, 10, 300]],
[[2, 20, 100], [2, 20, 200], [2, 20, 300]],
[[2, 30, 100], [2, 30, 200], [2, 30, 300]]
], [
[[3, 10, 100], [3, 10, 200], [3, 10, 300]],
[[3, 20, 100], [3, 20, 200], [3, 20, 300]],
[[3, 30, 100], [3, 30, 200], [3, 30, 300]]
]]

Вроде бы несложная задача. Однако я перерыл всю документацию, и не нашел никакого очевидного способа это сделать. Понятно, что можно вручную через несколько вложенных for, но это медленно (в реальности объем намного больше). Нужно в стиле NumPy.

Answer 1

[dmshar]

Можно еще вот так:

import itertools
x = [1, 2, 3]
y = [10, 20, 30]
z = [100, 200, 300]

ar=np.array(list(itertools.product(x,y,z))).reshape(len(x),len(y),len(z),-1)

Результат:

array([[[[  1,  10, 100],
         [  1,  10, 200],
         [  1,  10, 300]],

        [[  1,  20, 100],
         [  1,  20, 200],
         [  1,  20, 300]],

        [[  1,  30, 100],
         [  1,  30, 200],
         [  1,  30, 300]]],


       [[[  2,  10, 100],
         [  2,  10, 200],
         [  2,  10, 300]],

        [[  2,  20, 100],
         [  2,  20, 200],
         [  2,  20, 300]],

        [[  2,  30, 100],
         [  2,  30, 200],
         [  2,  30, 300]]],


       [[[  3,  10, 100],
         [  3,  10, 200],
         [  3,  10, 300]],

        [[  3,  20, 100],
         [  3,  20, 200],
         [  3,  20, 300]],

        [[  3,  30, 100],
         [  3,  30, 200],
         [  3,  30, 300]]]])

ar.shape
(3, 3, 3, 3)

Comments

[dom1n1k]

Спасибо, я не знал о такой функции, посмотрю. Выглядит очень похоже на правду, именно "перемножение" и подразумевалось.

Answer 2

[U235U235]

Можно использовать meshgrid и reshape. Что-то подобное:

import numpy as np
x = [1, 2, 3]
y = [10, 20, 30]
z = [100, 200, 300]
res1=np.array(np.meshgrid(x,y)).T.reshape(-1,3,2)
print(res1)
res2=np.array(np.meshgrid(x,y,z)).T.reshape(-1,3,3)
print(res2)

Comments

[dom1n1k]

Вроде бы работает, но я пока не совсем понял эту механику. Надо подумать.

[dmshar]

Если внимательно посмотреть на результат, запрашиваемый ТС, то видно, что он хочет получить массив
s=np.array([
[
[[1, 10, 100], [1, 10, 200], [1, 10, 300]],
[[1, 20, 100], [1, 20, 200], [1, 20, 300]],
[[1, 30, 100], [1, 30, 200], [1, 30, 300]]
],
[
[[2, 10, 100], [2, 10, 200], [2, 10, 300]],
[[2, 20, 100], [2, 20, 200], [2, 20, 300]],
[[2, 30, 100], [2, 30, 200], [2, 30, 300]]
],
[
[[3, 10, 100], [3, 10, 200], [3, 10, 300]],
[[3, 20, 100], [3, 20, 200], [3, 20, 300]],
[[3, 30, 100], [3, 30, 200], [3, 30, 300]]
]
])
А если проверить размерность этого массива s.shape то он будет (неожидано)
(3, 3, 3, 3)

Ваш же код дает

res2.shape
Out[88]: (9, 3, 3)

Не знаю, существенно ли это для ТС, но вот наличие отличия от задания - налицо.

[dom1n1k]

По идее, размерность итогового массива должна получиться (len(x), len(y), len(z), k), где k количество векторов