Search code examples
pythonnested-loopsprobabilitycombinatorics

How to remove a list from a nested list in python?


I'm exploring probability with the use of python and I want to resolve this kind of problem. I have 5 couples and from these couples, I want to extract each group of three, that doesn't contain persons from a married couple.

import itertools as it

married_couples = [["Mary","John"],["Homer","Marge"],["Beauty","Beast"],["Abigail","George"],["Marco","Luisa"]]

all_persons = []
for couples in married_couples:
    for person in couples:
        all_persons.append(person)

sample_space = list(it.combinations(all_persons,3))

# Better solution with generator:
sample_space = [s for t in it.combinations(married_couples, 3) for s in it.product(*t)]

Now I would like to proceed excluding from the sample space all results that contains people from the same married couple and I try:

correct_results = sample_space.copy()
for couple in married_couples:
    for res in sample_space:
        if couple[0] in res and couple[1] in res:
                if res in correct_results:
                    correct_results.remove(res)

Edit: I edited and put also the generator solution inside, so you can use this code for the purpose


Solution

  • The problem is in

    if couple[0] and couple[1] in res:
    

    because it tests not that the couple is in res, but that the first element of the couple is not null and the second is in res.

    You should use:

    if couple[0] in res and couple[1] in res: