Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1import numpy as np 

2 

3from pandas.core.dtypes.common import is_list_like 

4 

5import pandas.core.common as com 

6 

7 

8def cartesian_product(X): 

9 """ 

10 Numpy version of itertools.product. 

11 Sometimes faster (for large inputs)... 

12 

13 Parameters 

14 ---------- 

15 X : list-like of list-likes 

16 

17 Returns 

18 ------- 

19 product : list of ndarrays 

20 

21 Examples 

22 -------- 

23 >>> cartesian_product([list('ABC'), [1, 2]]) 

24 [array(['A', 'A', 'B', 'B', 'C', 'C'], dtype='|S1'), 

25 array([1, 2, 1, 2, 1, 2])] 

26 

27 See Also 

28 -------- 

29 itertools.product : Cartesian product of input iterables. Equivalent to 

30 nested for-loops. 

31 """ 

32 msg = "Input must be a list-like of list-likes" 

33 if not is_list_like(X): 

34 raise TypeError(msg) 

35 for x in X: 

36 if not is_list_like(x): 

37 raise TypeError(msg) 

38 

39 if len(X) == 0: 

40 return [] 

41 

42 lenX = np.fromiter((len(x) for x in X), dtype=np.intp) 

43 cumprodX = np.cumproduct(lenX) 

44 

45 a = np.roll(cumprodX, 1) 

46 a[0] = 1 

47 

48 if cumprodX[-1] != 0: 

49 b = cumprodX[-1] / cumprodX 

50 else: 

51 # if any factor is empty, the cartesian product is empty 

52 b = np.zeros_like(cumprodX) 

53 

54 return [ 

55 np.tile( 

56 np.repeat(np.asarray(com.values_from_object(x)), b[i]), np.product(a[i]) 

57 ) 

58 for i, x in enumerate(X) 

59 ]