3

There are k treatments and N total tests to distribute among the treatments, which is called a plan. For a fixed plan, I want to output in Python all the possible success sets.

Question:

For example, if doctors are testing headache medicine, if k=2 types of treatments (i.e. Aspirin and Ibuprofen) and N=3 total tests, one plan could be (1 test for Aspirin, 2 tests for Ibuprofen). For that plan, how do I output all possible combinations of 0-1 successful tests of Aspirin and 0-2 successful tests for Ibuprofen? One successful test means that when a patient with a headache is given Aspirin, the Aspirin cures their headache.

Please post an answer with python code, NOT a math answer.

Desired output is a list w/n a list that has [# successes for treatment 1, # successes of treatment 2]:

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

It would be great if yield could be used because the list above could be really long and I don't want to store the whole list in memory, which would increase computation time.

Below I have the code for enumerating all possible combinations of N balls in A boxes, which should be similar to creating all possible success sets I think, but I'm not sure how.

Code

#Return list of tuples of all possible plans (n1,..,nk), where N = total # of tests = balls, K = # of treatments = boxes
#Code: Glyph, http://stackoverflow.com/questions/996004/enumeration-of-combinations-of-n-balls-in-a-boxes
def ballsAndBoxes(balls, boxes, boxIndex=0, sumThusFar=0):
    if boxIndex < (boxes - 1):
        for counter in range(balls + 1 - sumThusFar):
            for rest in ballsAndBoxes(balls, boxes,
                                      boxIndex + 1,
                                      sumThusFar + counter):
                yield (counter,) + rest
    else:
        yield (balls - sumThusFar,)
14wml
  • 4,048
  • 11
  • 49
  • 97

2 Answers2

1

Generating the plans is a partition problem, but generating the success sets for a given plan only requires generating the Cartesian product of a set of ranges.

from itertools import product

def success_sets(plan):
    return product(*map(lambda n: range(n + 1), plan))

plan = [1, 2]
for s in success_sets(plan):
    print(s)
# (0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)

Since itertools.product returns a generator, the entire list will not be stored in memory as requested.

Jared Goguen
  • 8,772
  • 2
  • 18
  • 36
0

I am not sure exactly what you're trying to achieve. But combinations can be generated using itertools.

from itertools import combinations
    #You can add an extra loop for all treatments
    for j in range(1, N): #N is number of tests
        for i in combinations(tests, r = j):
            indexes = set(i)
            df_cur = tests[indexes] #for tests i am using a pandas df
            if :# condition for success
                #actions
            else:
                #other_actions
Abhishek Menon
  • 753
  • 4
  • 15
  • Hi, I updated my question with a real life example, so hopefully my question is more clear. I don't think this answers my question because I don't need a condition for success. For example, for a fixed plan (1 test for Aspirin, 2 test for Ibuprofen), I'm trying to output the possible combination of successful tests, which are (0 successful tests for Aspirin, 0 successful tests for Ibuprofen), (0 successful tests for Aspirin, 1 successful test for Ibuprofen), (0, 2), (1, 0), (1,1), (1, 2). Does that make sense? – 14wml Apr 20 '17 at 17:28