Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import annotations
3from typing import TYPE_CHECKING
5import numpy as np
7from pandas.core.dtypes.common import is_list_like
9if TYPE_CHECKING:
10 from pandas._typing import NumpyIndexT
13def cartesian_product(X) -> list[np.ndarray]:
14 """
15 Numpy version of itertools.product.
16 Sometimes faster (for large inputs)...
18 Parameters
19 ----------
20 X : list-like of list-likes
22 Returns
23 -------
24 product : list of ndarrays
26 Examples
27 --------
28 >>> cartesian_product([list('ABC'), [1, 2]])
29 [array(['A', 'A', 'B', 'B', 'C', 'C'], dtype='<U1'), array([1, 2, 1, 2, 1, 2])]
31 See Also
32 --------
33 itertools.product : Cartesian product of input iterables. Equivalent to
34 nested for-loops.
35 """
36 msg = "Input must be a list-like of list-likes"
37 if not is_list_like(X):
38 raise TypeError(msg)
39 for x in X:
40 if not is_list_like(x):
41 raise TypeError(msg)
43 if len(X) == 0:
44 return []
46 lenX = np.fromiter((len(x) for x in X), dtype=np.intp)
47 cumprodX = np.cumprod(lenX)
49 if np.any(cumprodX < 0):
50 raise ValueError("Product space too large to allocate arrays!")
52 a = np.roll(cumprodX, 1)
53 a[0] = 1
55 if cumprodX[-1] != 0:
56 b = cumprodX[-1] / cumprodX
57 else:
58 # if any factor is empty, the cartesian product is empty
59 b = np.zeros_like(cumprodX)
61 # error: Argument of type "int_" cannot be assigned to parameter "num" of
62 # type "int" in function "tile_compat"
63 return [
64 tile_compat(
65 np.repeat(x, b[i]),
66 np.prod(a[i]),
67 )
68 for i, x in enumerate(X)
69 ]
72def tile_compat(arr: NumpyIndexT, num: int) -> NumpyIndexT:
73 """
74 Index compat for np.tile.
76 Notes
77 -----
78 Does not support multi-dimensional `num`.
79 """
80 if isinstance(arr, np.ndarray):
81 return np.tile(arr, num)
83 # Otherwise we have an Index
84 taker = np.tile(np.arange(len(arr)), num)
85 return arr.take(taker)