Search code examples
pythonlistpython-itertools

Keeping track of which subset is computed to avoid mixing the results


tricky name and unclear name...

My issue is the following:

import itertools

t0 = 0
tf = 1000000
# inputs_to_compute = list-like of size 2 to 6 of objects

results = [[] for i in range(len(inputs_to_compute))]

for subset in itertools.combinations(inputs_to_compute, 2):
    r1, r2 = compute(subset[0], subset[1], t0, tf)
    results[inputs_to_compute.index(subset[0])] += list(r1)
    results[inputs_to_compute.index(subset[1])] += list(r2)

This code is creating as many results lists as there is an input. Each input is actually associated to a list. Then the computation is performed 2 by 2 (on each subset) and the result is added to the corresponding list.

It works well as long as there is no repetition in the inputs because the method index returns the first occurrence of the element. How could I implement this differently (and efficiently, performance is one of the main issue I have) in a way that manage duplicates?

Dummy example:

import itertools

def compute(x, y):
    return (x + y, x - y)

inputs_to_compute = [1, 1, 3]

results = [[] for i in range(len(inputs_to_compute))]

for subset in itertools.combinations(inputs_to_compute, 2):
    r1, r2 = compute(subset[0], subset[1])
    results[inputs_to_compute.index(subset[0])].append(r1)
    results[inputs_to_compute.index(subset[1])].append(r2)

Output:

[[2, 0, 4, 4], [], [-2, -2]]

Expected output:

# Iteration (1, 1): r1 = 2, r2 = 0
results = [[2], [0], []]
# Iteration (1, 3): r1 = 4, r2 = -2
results = [[2, 4], [0], [-2]]
# Iteration (1, 3): r1 = 4, r2 = -2
results = [[2, 4], [0, 4], [-2, -2]]

Solution

  • for subset_with_indices in itertools.combinations(enumerate(inputs_to_compute), 2):
        i1,x1 = subset_with_indices[0]
        i2,x2 = subset_with_indices[1]
        r1, r2 = compute(x1, x2)
        results[i1].append(r1)
        results[i2].append(r2)