27

The result created by Python's itertools.combinations() is the combinations of numbers. For example:

a = [7, 5, 5, 4]
b = list(itertools.combinations(a, 2))

# b = [(7, 5), (7, 5), (7, 4), (5, 5), (5, 4), (5, 4)]

But I would like to also obtain the indices of the combinations such as:

index = [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]

How can I do it?

Randy Tang
  • 4,283
  • 12
  • 54
  • 112

3 Answers3

28

You can use enumerate:

>>> a = [7, 5, 5, 4]
>>> list(itertools.combinations(enumerate(a), 2))
[((0, 7), (1, 5)), ((0, 7), (2, 5)), ((0, 7), (3, 4)), ((1, 5), (2, 5)), ((1, 5), (3, 4)), ((2, 5), (3, 4))]
>>> b = list((i,j) for ((i,_),(j,_)) in itertools.combinations(enumerate(a), 2))
>>> b
[(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]
fredtantini
  • 15,966
  • 8
  • 49
  • 55
6

You can use range to get order of indexes that combinations produce.

>>> list(itertools.combinations(range(3), 2))
[(0, 1), (0, 2), (1, 2)]

So you can use len(a):

>>> list(itertools.combinations(range(len(a)), 2))
[(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]
Tommaso Di Noto
  • 1,208
  • 1
  • 13
  • 24
Ardalan Hosseini
  • 61
  • 1
  • 1
0
a = [7, 5, 5, 4]
n_combinations = 2

np.array(list(itertools.combinations(enumerate(a), n_combinations)))[...,0]

More extendable solution. Change of n_combinations does not require change of code, as it does in the example from @fredtantini. It is basically his solution expressed with numpy slicing.

Petr Dvořáček
  • 123
  • 6