Build all combinations of success sets in Python
There are k treatments and N complete tests to distribute to a treatment called a plan. For a fixed plan, I want to output all possible success options in Python.
Question:
For example, if doctors are testing a headache medicine, if k = 2 treatments (i.e., aspirin and ibuprofen) and N = 3 general tests, one plan might be (1 aspirin test, 2 ibuprofen tests). For this plan, how can I output all possible combinations of 0-1 successful aspirin tests and 0-2 successful tests for Ibuprofen? One successful test means that when an aspirin is prescribed to a patient with a headache, the aspirin cures their headache.
Please post an answer with python code , not a mathematical answer.
Desired output is a list w / n list which has [# of success for treatment 1, # of success for treatment 2]:
[[0.0], [0.1], [0.2], [1.0], [1.1], [1.2]]
It would be great if it yield
could be used because the above list can be very long and I don't want to keep the entire list in memory, which would increase the computation time.
Below I have the code to list all the possible combinations of N balls in A blocks, which should be like creating all possible sets of successes 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,)
source to share
Creating plans is a sectional issue, but creating success sets for a given plan only requires creating the Cartesian product of a range set.
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 the itertools.product
generator returns, the entire list will not be stored in memory upon request.
source to share
I don't know exactly what you are 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
source to share