Найти все комбинации, состоящие из одной строки в нескольких массивах

Question

• У меня есть несколько двумерных массивов.
• Я хочу найти все комбинации, состоящие из одной строки в каждом массиве.

Пример:

Допустим, в массиве A есть строки a0, a1, a2. Допустим, массив B имеет строки b0, b1

Шесть комбинаций:

a0-b0, a0-b1, a1-b0, a1-b1, a2-b0, a2-b1

Тире представляет собой конкатенацию (np.hstack)

Как сделать это быстро для произвольного количества массивов (скажем, A, B, C, ...)?

Быстрый метод для 2 массивов: Комбинация всех рядов в двух пустых массивах

Результат кода для объединения 3 массивов:

# input arrays:
[[0 1 2]]
[[ 5  6  7  8  9]
 [10 11 12 13 14]
 [15 16 17 18 19]]
[[17 18 19 20 21 22 23]
 [24 25 26 27 28 29 30]
 [31 32 33 34 35 36 37]]
# output:
[[ 0  1  2  5  6  7  8  9 17 18 19 20 21 22 23]
 [ 0  1  2  5  6  7  8  9 24 25 26 27 28 29 30]
 [ 0  1  2  5  6  7  8  9 31 32 33 34 35 36 37]
 [ 0  1  2 10 11 12 13 14 17 18 19 20 21 22 23]
 [ 0  1  2 10 11 12 13 14 24 25 26 27 28 29 30]
 [ 0  1  2 10 11 12 13 14 31 32 33 34 35 36 37]
 [ 0  1  2 15 16 17 18 19 17 18 19 20 21 22 23]
 [ 0  1  2 15 16 17 18 19 24 25 26 27 28 29 30]
 [ 0  1  2 15 16 17 18 19 31 32 33 34 35 36 37]]

Код:

import numpy as np

def form_combinations(xs):
    tot_size = np.sum([x.shape[1] for x in xs])
    n_rows = [x.shape[0] for x in xs]
    out = np.empty(n_rows + [tot_size])
    n_cols = [x.shape[1] for x in xs]
    cs = np.cumsum([0] + n_cols)
    n = np.newaxis
    out[:, :, :, cs[0]:cs[1]] = xs[0][:, n, n, :]
    out[:, :, :, cs[1]:cs[2]] = xs[1][n, :, n, :]
    out[:, :, :, cs[2]:cs[3]] = xs[2][n, n, :, :]
    out = out.reshape(-1, tot_size)
    return out


def main():
    xs = [
        np.arange(3)[np.newaxis, :],
        np.arange(5, 20).reshape(3, 5),
        np.arange(17, 38).reshape(3, 7)
    ]
    print(xs)

    out = form_combinations(xs)
    print(out)

main()

Answer 1

Для произвольного числа массивов, инкрементальный подход:

def combine(xs):
    comb=np.array([[]],int)
    for array in xs:
        left  = repeat(comb,len(array),axis=0)
        right = vstack([array]*len(comb))
        comb  = hstack((left,right))
    return comb

Answer 2

Адаптировано с https://stackoverflow.com/a/49445693/7207392

import numpy as np
import operator as op
import itertools as it

def cartesian_product_pp(arrays, out=None):
    la = len(arrays)
    h, w = zip(*map(op.attrgetter('shape'), arrays))
    w = np.fromiter(it.chain([0], w), int, la+ 1)
    W = w.cumsum()
    h = *h, W[la]
    dtype = np.result_type(*arrays)
    arr = np.empty(h, dtype=dtype)
    arrs = *it.accumulate(it.chain((arr,), it.repeat(0, la-1)), np.ndarray.__getitem__),
    idx = slice(None), *it.repeat(None, la-1)
    for i in range(la-1, 0, -1):
        arrs[i][..., W[i]:W[i+1]] = arrays[i][idx[:la-i]]
        arrs[i-1][1:] = arrs[i]
    arr[..., W[0]:W[1]] = arrays[0][idx]
    return arr.reshape(-1, W[la])

# example
a = np.r_[:3].reshape(1, 3)
b = np.r_[5:20].reshape(3, 5)
c = np.r_[17:38].reshape(3, 7)
p = cartesian_product_pp([a, b, c])

Выход:

>>> p
array([[ 0,  1,  2,  5,  6,  7,  8,  9, 17, 18, 19, 20, 21, 22, 23],
       [ 0,  1,  2,  5,  6,  7,  8,  9, 24, 25, 26, 27, 28, 29, 30],
       [ 0,  1,  2,  5,  6,  7,  8,  9, 31, 32, 33, 34, 35, 36, 37],
       [ 0,  1,  2, 10, 11, 12, 13, 14, 17, 18, 19, 20, 21, 22, 23],
       [ 0,  1,  2, 10, 11, 12, 13, 14, 24, 25, 26, 27, 28, 29, 30],
       [ 0,  1,  2, 10, 11, 12, 13, 14, 31, 32, 33, 34, 35, 36, 37],
       [ 0,  1,  2, 15, 16, 17, 18, 19, 17, 18, 19, 20, 21, 22, 23],
       [ 0,  1,  2, 15, 16, 17, 18, 19, 24, 25, 26, 27, 28, 29, 30],
       [ 0,  1,  2, 15, 16, 17, 18, 19, 31, 32, 33, 34, 35, 36, 37]])

Время для этого, подходы @ B.M. и @ Bazingaa:

>>> timeit(lambda: cartesian_product_pp([a,b,c]), number=1000)*1000
15.173833002336323
>>> timeit(lambda: combine([a,b,c]), number=1000)*1000
31.1394709860906
>>> timeit(lambda: np.array([np.hstack((i, j, k)) for i in a for j in b for k in c]), number=1000)*1000
51.15771805867553

Answer 3

Одним из способов является использование списка, где вы просто зацикливаетесь на всех трех массивах, и используете hstack, чтобы расположить их горизонтально.

np.array([np.hstack((i, j, k)) for i in a for j in b for k in c])

# array([[ 0,  1,  2,  5,  6,  7,  8,  9, 17, 18, 19, 20, 21, 22, 23],
#        [ 0,  1,  2,  5,  6,  7,  8,  9, 24, 25, 26, 27, 28, 29, 30],
#        [ 0,  1,  2,  5,  6,  7,  8,  9, 31, 32, 33, 34, 35, 36, 37],
#        [ 0,  1,  2, 10, 11, 12, 13, 14, 17, 18, 19, 20, 21, 22, 23],
#        [ 0,  1,  2, 10, 11, 12, 13, 14, 24, 25, 26, 27, 28, 29, 30],
#        [ 0,  1,  2, 10, 11, 12, 13, 14, 31, 32, 33, 34, 35, 36, 37],
#        [ 0,  1,  2, 15, 16, 17, 18, 19, 17, 18, 19, 20, 21, 22, 23],
#        [ 0,  1,  2, 15, 16, 17, 18, 19, 24, 25, 26, 27, 28, 29, 30],
#        [ 0,  1,  2, 15, 16, 17, 18, 19, 31, 32, 33, 34, 35, 36, 37]])

Задержка

%timeit np.array([np.hstack((i, j, k)) for i in a for j in b for k in c])

# 55.1 µs ± 2.61 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)