CODEWITHSUNDEEP
python-3.x
ratorx
ratorx

Reputation: 357

Efficient way to iterate over sublists of a list grouped by duplicates

The title may be unclear, but I can't think of a better way to phrase it. I mean a "list grouped by duplicates" as:

[2, 2, 1]

where each number in this list represents the number of duplicates of an element linked to the ith position of the list. For example, the above list could represent:

[0, 0, 1, 1, 2]

or

["a", "a", "b", "b", "c"]

or anything like that.

What would be an efficient way to generate/iterate over all the sublists of such a list?

As an example, the iterator of the above list should produce (in some order):

[[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1], [1, 1, 0], [1, 0, 1], [0, 1, 1], [2, 0, 0], [0, 2, 0], [1, 1, 1], [2, 1, 0], [2, 0, 1], [0, 2, 1], [2, 1, 1], [2, 2, 0], [1, 2, 1], [2, 2, 1]]

I think the above list is complete, if it's not, please point it out. Regardless, I hope you have an idea of what I'm trying to do. Thanks in advance

Upvotes: 1

Views: 29

Answers (1)

Apollys supports Monica
Apollys supports Monica

Reputation: 3258

For clarity, first we can generate the list of options for possible counts of each element, then take the cartesian product to get the desired result.

import itertools # for itertools.product
a = [2, 2, 1]
options = [list(range(x+1)) for x in a]
result = list(itertools.product(*options))
print(result)

[(0, 0, 0),
 (0, 0, 1),
 (0, 1, 0),
 (0, 1, 1),
 (0, 2, 0),
 (0, 2, 1),
 (1, 0, 0),
 (1, 0, 1),
 (1, 1, 0),
 (1, 1, 1),
 (1, 2, 0),
 (1, 2, 1),
 (2, 0, 0),
 (2, 0, 1),
 (2, 1, 0),
 (2, 1, 1),
 (2, 2, 0),
 (2, 2, 1)]

Note that the * is an argument unpacking operator here, which allows the itertools.product function to take the product of all the (arbitrarily many) sublists we gave it in the options list.

Upvotes: 3

Related Questions