3

What's the best way to get a pandas dataframe with two columns that provide all possible combinations of numbers in two lists?

I've used itertools to get a list of lists with the combinations, but can't work out how to get it into a dataframe

Here's what I've tried

import itertools
import pandas as pd

income_range = np.linspace(1,3,3)
costs_range = np.linspace(1, 3, 3)
all_combinations = [list(zip(each_permutation, costs_range)) for each_permutation in itertools.permutations(income_range, len(costs_range))]
df = pd.DataFrame(all_combinations)

This is the output that I get

|    | 0          | 1          | 2          |
|---:|:-----------|:-----------|:-----------|
|  0 | (1.0, 1.0) | (2.0, 2.0) | (3.0, 3.0) |
|  1 | (1.0, 1.0) | (3.0, 2.0) | (2.0, 3.0) |
|  2 | (2.0, 1.0) | (1.0, 2.0) | (3.0, 3.0) |
|  3 | (2.0, 1.0) | (3.0, 2.0) | (1.0, 3.0) |
|  4 | (3.0, 1.0) | (1.0, 2.0) | (2.0, 3.0) |
|  5 | (3.0, 1.0) | (2.0, 2.0) | (1.0, 3.0) |

This is the output that I want

| Income | Costs |
| ------ | ----- |
| 1      | 1     |
| 1      | 2     |
| 1      | 3     |
| 2      | 1     |
| 2      | 2     |
| 2      | 3     |
| 3      | 1     |
| 3      | 2     |
| 3      | 3     |
Vadim Kotov
  • 8,084
  • 8
  • 48
  • 62
yahop19531
  • 193
  • 1
  • 11

2 Answers2

3

Use itertools.product here:

income_range = np.linspace(1,3,3)
costs_range = np.linspace(1, 3, 3)

all_combinations = list(itertools.product(income_range, costs_range))
df = pd.DataFrame(all_combinations, columns=['Income','Costs'])
print (df)
   Income  Costs
0     1.0    1.0
1     1.0    2.0
2     1.0    3.0
3     2.0    1.0
4     2.0    2.0
5     2.0    3.0
6     3.0    1.0
7     3.0    2.0
8     3.0    3.0
jezrael
  • 822,522
  • 95
  • 1,334
  • 1,252
1

You could use Meshgrid to solve this:

result = np.reshape(np.vstack(np.meshgrid(income_range, costs_range)), (-1, 2))
pd.DataFrame(result)

    0   1
0   1.0 2.0
1   3.0 1.0
2   2.0 3.0
3   1.0 2.0
4   3.0 1.0
5   1.0 1.0
6   2.0 2.0
7   2.0 3.0
8   3.0 3.0

You can have a look here for possible optimizations. Do note that this excels with numbers; for other data types, especially strings, I believe itertools' product function is apt (@jezrael's solution)

sammywemmy
  • 27,093
  • 4
  • 17
  • 31