Um canônico cartesian_product(quase)
Existem muitas abordagens para esse problema com propriedades diferentes. Alguns são mais rápidos que outros, e outros são de uso geral. Após muitos testes e ajustes, descobri que a função a seguir, que calcula uma dimensão n cartesian_product, é mais rápida que a maioria das outras para muitas entradas. Para um par de abordagens um pouco mais complexas, mas ainda mais rápidas em muitos casos, veja a resposta de Paul Panzer .
Dada essa resposta, essa não é mais a implementação mais rápida do produto cartesiano em numpyque estou ciente. No entanto, acho que sua simplicidade continuará sendo uma referência útil para melhorias futuras:
def cartesian_product(*arrays):
la = len(arrays)
dtype = numpy.result_type(*arrays)
arr = numpy.empty([len(a) for a in arrays] + [la], dtype=dtype)
for i, a in enumerate(numpy.ix_(*arrays)):
arr[...,i] = a
return arr.reshape(-1, la)
Vale ressaltar que essa função usa de ix_maneira incomum; enquanto o uso documentado de ix_é gerar índices em uma matriz, acontece que matrizes com a mesma forma podem ser usadas para atribuição transmitida. Muito obrigado a mgilson , que me inspirou a tentar usar ix_esse caminho, e a unutbu , que forneceu alguns comentários extremamente úteis sobre essa resposta, incluindo a sugestão de uso numpy.result_type.
Alternativas notáveis
Às vezes, é mais rápido gravar blocos de memória contíguos na ordem do Fortran. Essa é a base dessa alternativa, cartesian_product_transposeque se mostrou mais rápida em alguns hardwares do que cartesian_product(veja abaixo). No entanto, a resposta de Paul Panzer, que usa o mesmo princípio, é ainda mais rápida. Ainda assim, incluo isso aqui para leitores interessados:
def cartesian_product_transpose(*arrays):
broadcastable = numpy.ix_(*arrays)
broadcasted = numpy.broadcast_arrays(*broadcastable)
rows, cols = numpy.prod(broadcasted[0].shape), len(broadcasted)
dtype = numpy.result_type(*arrays)
out = numpy.empty(rows * cols, dtype=dtype)
start, end = 0, rows
for a in broadcasted:
out[start:end] = a.reshape(-1)
start, end = end, end + rows
return out.reshape(cols, rows).T
Depois de entender a abordagem de Panzer, escrevi uma nova versão quase tão rápida quanto a dele e quase tão simples quanto cartesian_product:
def cartesian_product_simple_transpose(arrays):
la = len(arrays)
dtype = numpy.result_type(*arrays)
arr = numpy.empty([la] + [len(a) for a in arrays], dtype=dtype)
for i, a in enumerate(numpy.ix_(*arrays)):
arr[i, ...] = a
return arr.reshape(la, -1).T
Parece haver alguma sobrecarga de tempo constante que a torna mais lenta que a do Panzer para pequenas entradas. Mas para entradas maiores, em todos os testes que eu executei, ele executa tão bem quanto sua implementação mais rápida ( cartesian_product_transpose_pp).
Nas seções a seguir, incluo alguns testes de outras alternativas. Agora eles estão um pouco desatualizados, mas, em vez de um esforço duplicado, decidi deixá-los aqui fora de interesse histórico. Para testes atualizados, consulte a resposta de Panzer, bem como a de Nico Schlömer .
Testes contra alternativas
Aqui está uma bateria de testes que mostram o aumento de desempenho que algumas dessas funções fornecem em relação a várias alternativas. Todos os testes mostrados aqui foram realizados em uma máquina quad-core, executando o Mac OS 10.12.5, Python 3.6.1 e numpy1.12.1. Sabe-se que variações no hardware e software produzem resultados diferentes, portanto, YMMV. Execute esses testes para ter certeza!
Definições:
import numpy
import itertools
from functools import reduce
### Two-dimensional products ###
def repeat_product(x, y):
return numpy.transpose([numpy.tile(x, len(y)),
numpy.repeat(y, len(x))])
def dstack_product(x, y):
return numpy.dstack(numpy.meshgrid(x, y)).reshape(-1, 2)
### Generalized N-dimensional products ###
def cartesian_product(*arrays):
la = len(arrays)
dtype = numpy.result_type(*arrays)
arr = numpy.empty([len(a) for a in arrays] + [la], dtype=dtype)
for i, a in enumerate(numpy.ix_(*arrays)):
arr[...,i] = a
return arr.reshape(-1, la)
def cartesian_product_transpose(*arrays):
broadcastable = numpy.ix_(*arrays)
broadcasted = numpy.broadcast_arrays(*broadcastable)
rows, cols = numpy.prod(broadcasted[0].shape), len(broadcasted)
dtype = numpy.result_type(*arrays)
out = numpy.empty(rows * cols, dtype=dtype)
start, end = 0, rows
for a in broadcasted:
out[start:end] = a.reshape(-1)
start, end = end, end + rows
return out.reshape(cols, rows).T
# from https://stackoverflow.com/a/1235363/577088
def cartesian_product_recursive(*arrays, out=None):
arrays = [numpy.asarray(x) for x in arrays]
dtype = arrays[0].dtype
n = numpy.prod([x.size for x in arrays])
if out is None:
out = numpy.zeros([n, len(arrays)], dtype=dtype)
m = n // arrays[0].size
out[:,0] = numpy.repeat(arrays[0], m)
if arrays[1:]:
cartesian_product_recursive(arrays[1:], out=out[0:m,1:])
for j in range(1, arrays[0].size):
out[j*m:(j+1)*m,1:] = out[0:m,1:]
return out
def cartesian_product_itertools(*arrays):
return numpy.array(list(itertools.product(*arrays)))
### Test code ###
name_func = [('repeat_product',
repeat_product),
('dstack_product',
dstack_product),
('cartesian_product',
cartesian_product),
('cartesian_product_transpose',
cartesian_product_transpose),
('cartesian_product_recursive',
cartesian_product_recursive),
('cartesian_product_itertools',
cartesian_product_itertools)]
def test(in_arrays, test_funcs):
global func
global arrays
arrays = in_arrays
for name, func in test_funcs:
print('{}:'.format(name))
%timeit func(*arrays)
def test_all(*in_arrays):
test(in_arrays, name_func)
# `cartesian_product_recursive` throws an
# unexpected error when used on more than
# two input arrays, so for now I've removed
# it from these tests.
def test_cartesian(*in_arrays):
test(in_arrays, name_func[2:4] + name_func[-1:])
x10 = [numpy.arange(10)]
x50 = [numpy.arange(50)]
x100 = [numpy.arange(100)]
x500 = [numpy.arange(500)]
x1000 = [numpy.arange(1000)]
Resultado dos testes:
In [2]: test_all(*(x100 * 2))
repeat_product:
67.5 µs ± 633 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
dstack_product:
67.7 µs ± 1.09 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
cartesian_product:
33.4 µs ± 558 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
cartesian_product_transpose:
67.7 µs ± 932 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
cartesian_product_recursive:
215 µs ± 6.01 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
cartesian_product_itertools:
3.65 ms ± 38.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [3]: test_all(*(x500 * 2))
repeat_product:
1.31 ms ± 9.28 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
dstack_product:
1.27 ms ± 7.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
cartesian_product:
375 µs ± 4.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
cartesian_product_transpose:
488 µs ± 8.88 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
cartesian_product_recursive:
2.21 ms ± 38.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
cartesian_product_itertools:
105 ms ± 1.17 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [4]: test_all(*(x1000 * 2))
repeat_product:
10.2 ms ± 132 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
dstack_product:
12 ms ± 120 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
cartesian_product:
4.75 ms ± 57.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
cartesian_product_transpose:
7.76 ms ± 52.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
cartesian_product_recursive:
13 ms ± 209 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
cartesian_product_itertools:
422 ms ± 7.77 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
Em todos os casos, cartesian_productconforme definido no início desta resposta, é mais rápido.
Para aquelas funções que aceitam um número arbitrário de matrizes de entrada, vale a pena verificar o desempenho len(arrays) > 2também. (Até que eu possa determinar por que causa cartesian_product_recursiveum erro nesse caso, eu o removi desses testes.)
In [5]: test_cartesian(*(x100 * 3))
cartesian_product:
8.8 ms ± 138 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
cartesian_product_transpose:
7.87 ms ± 91.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
cartesian_product_itertools:
518 ms ± 5.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [6]: test_cartesian(*(x50 * 4))
cartesian_product:
169 ms ± 5.1 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
cartesian_product_transpose:
184 ms ± 4.32 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
cartesian_product_itertools:
3.69 s ± 73.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [7]: test_cartesian(*(x10 * 6))
cartesian_product:
26.5 ms ± 449 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
cartesian_product_transpose:
16 ms ± 133 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
cartesian_product_itertools:
728 ms ± 16 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [8]: test_cartesian(*(x10 * 7))
cartesian_product:
650 ms ± 8.14 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
cartesian_product_transpose:
518 ms ± 7.09 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
cartesian_product_itertools:
8.13 s ± 122 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
Como esses testes mostram, cartesian_productpermanece competitivo até que o número de matrizes de entrada ultrapasse (aproximadamente) quatro. Depois disso, cartesian_product_transposetem uma ligeira vantagem.
Vale a pena reiterar que usuários com outros hardwares e sistemas operacionais podem ter resultados diferentes. Por exemplo, o unutbu reporta os seguintes resultados para esses testes usando o Ubuntu 14.04, Python 3.4.3 e numpy1.14.0.dev0 + b7050a9:
>>> %timeit cartesian_product_transpose(x500, y500)
1000 loops, best of 3: 682 µs per loop
>>> %timeit cartesian_product(x500, y500)
1000 loops, best of 3: 1.55 ms per loop
Abaixo, mostro alguns detalhes sobre os testes anteriores que realizei nesse sentido. O desempenho relativo dessas abordagens mudou ao longo do tempo, para diferentes hardwares e diferentes versões do Python e numpy. Embora não seja imediatamente útil para pessoas que usam versões atualizadas numpy, ele ilustra como as coisas mudaram desde a primeira versão desta resposta.
Uma alternativa simples: meshgrid+dstack
A resposta atualmente aceita usa tilee repeatpara transmitir duas matrizes juntas. Mas a meshgridfunção faz praticamente a mesma coisa. Aqui está a saída tilee repeatantes de ser passada para transposição:
In [1]: import numpy
In [2]: x = numpy.array([1,2,3])
...: y = numpy.array([4,5])
...:
In [3]: [numpy.tile(x, len(y)), numpy.repeat(y, len(x))]
Out[3]: [array([1, 2, 3, 1, 2, 3]), array([4, 4, 4, 5, 5, 5])]
E aqui está a saída de meshgrid:
In [4]: numpy.meshgrid(x, y)
Out[4]:
[array([[1, 2, 3],
[1, 2, 3]]), array([[4, 4, 4],
[5, 5, 5]])]
Como você pode ver, é quase idêntico. Precisamos apenas remodelar o resultado para obter exatamente o mesmo resultado.
In [5]: xt, xr = numpy.meshgrid(x, y)
...: [xt.ravel(), xr.ravel()]
Out[5]: [array([1, 2, 3, 1, 2, 3]), array([4, 4, 4, 5, 5, 5])]
Em vez de remodelar neste ponto, porém, poderíamos passar a saída meshgridpara dstacke remodelar depois, o que economiza algum trabalho:
In [6]: numpy.dstack(numpy.meshgrid(x, y)).reshape(-1, 2)
Out[6]:
array([[1, 4],
[2, 4],
[3, 4],
[1, 5],
[2, 5],
[3, 5]])
Ao contrário da afirmação deste comentário , não vi nenhuma evidência de que entradas diferentes produzirão saídas de formas diferentes e, como demonstrado acima, elas fazem coisas muito semelhantes, portanto seria muito estranho se o fizessem. Entre em contato se você encontrar um contra-exemplo.
Testando meshgrid+ dstackvs. repeat+transpose
O desempenho relativo dessas duas abordagens mudou ao longo do tempo. Em uma versão anterior do Python (2.7), o resultado usando meshgrid+ dstackfoi notavelmente mais rápido para pequenas entradas. (Observe que esses testes são de uma versão antiga desta resposta.) Definições:
>>> def repeat_product(x, y):
... return numpy.transpose([numpy.tile(x, len(y)),
numpy.repeat(y, len(x))])
...
>>> def dstack_product(x, y):
... return numpy.dstack(numpy.meshgrid(x, y)).reshape(-1, 2)
...
Para entrada de tamanho médio, vi uma aceleração significativa. Mas tentei novamente esses testes com versões mais recentes do Python (3.6.1) e numpy(1.12.1), em uma máquina mais nova. As duas abordagens são quase idênticas agora.
Teste antigo
>>> x, y = numpy.arange(500), numpy.arange(500)
>>> %timeit repeat_product(x, y)
10 loops, best of 3: 62 ms per loop
>>> %timeit dstack_product(x, y)
100 loops, best of 3: 12.2 ms per loop
Novo teste
In [7]: x, y = numpy.arange(500), numpy.arange(500)
In [8]: %timeit repeat_product(x, y)
1.32 ms ± 24.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [9]: %timeit dstack_product(x, y)
1.26 ms ± 8.47 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
Como sempre, YMMV, mas isso sugere que nas versões recentes do Python e numpy, elas são intercambiáveis.
Funções generalizadas do produto
Em geral, podemos esperar que o uso de funções internas seja mais rápido para entradas pequenas, enquanto que para entradas grandes, uma função criada para fins específicos pode ser mais rápida. Além disso, para um produto n-dimensional generalizado, tilee repeatnão ajudará, porque eles não possuem análogos claros de alta dimensão. Portanto, vale a pena investigar o comportamento das funções criadas especificamente para esse fim.
A maioria dos testes relevantes aparece no início desta resposta, mas aqui estão alguns dos testes realizados em versões anteriores do Python e numpypara comparação.
A cartesianfunção definida em outra resposta costumava ter um desempenho muito bom para entradas maiores. (É o mesmo que a função chamada cartesian_product_recursiveacima.) A fim de comparar cartesiana dstack_prodct, usamos apenas duas dimensões.
Aqui, novamente, o teste antigo mostrou uma diferença significativa, enquanto o novo teste mostra quase nenhum.
Teste antigo
>>> x, y = numpy.arange(1000), numpy.arange(1000)
>>> %timeit cartesian([x, y])
10 loops, best of 3: 25.4 ms per loop
>>> %timeit dstack_product(x, y)
10 loops, best of 3: 66.6 ms per loop
Novo teste
In [10]: x, y = numpy.arange(1000), numpy.arange(1000)
In [11]: %timeit cartesian([x, y])
12.1 ms ± 199 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [12]: %timeit dstack_product(x, y)
12.7 ms ± 334 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
Como antes, dstack_productainda bate cartesianem escalas menores.
Novo teste ( teste antigo redundante não mostrado )
In [13]: x, y = numpy.arange(100), numpy.arange(100)
In [14]: %timeit cartesian([x, y])
215 µs ± 4.75 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [15]: %timeit dstack_product(x, y)
65.7 µs ± 1.15 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
Acho que essas distinções são interessantes e merecem ser registradas; mas eles são acadêmicos no final. Como mostraram os testes no início desta resposta, todas essas versões são quase sempre mais lentas que cartesian_product, definidas no início desta resposta - o que é um pouco mais lento que as implementações mais rápidas entre as respostas a essa pergunta.
