1from itertools import combinations, permutations
 2from copy import deepcopy
 3
 4def is_everyone_comfy(students):
 5    """Check if everyone is comfy in their current car."""
 6    for student in students:
 7        if not student.is_comfy():
 8            return False
 9            
10    return True
11            
12# implementation from StackOverflow - same as more_itertools.set_partitions
13def set_partitions_helper(collection):
14    """Returns a generator that yields set partitions of a collection."""
15    if len(collection) == 1:
16        yield [collection]
17        return
18    
19    first = collection[0]
20    for smaller in set_partitions_helper(collection[1:]):
21        for n, subset in enumerate(smaller):
22            yield smaller[:n] + [[ first ] + subset]  + smaller[n+1:]
23        yield [[first]] + smaller
24        
25
26def organize(cars, students):
27    """Organize the students in cars."""
28    if len(students) == 0:
29        return True
30    
31    if len(cars) * 4 < len(students):
32        return False
33    
34    students_configurations_generator = set_partitions_helper(students)
35
36    cars_configurations = [j for i in range(1, len(cars) + 1) for j in list(combinations(cars, i))]
37    
38    try:
39       while True:
40            students_configuration = next(students_configurations_generator)
41            if  not all(len(i) <= 4 for i in students_configuration):
42                continue
43            for cars_conf in cars_configurations:
44                if len(cars_conf) != len(students_configuration):
45                    continue
46
47                for cars_perm in permutations(cars_conf):
48                    cars_copy = deepcopy(cars_perm)
49                    students_copy = deepcopy(students_configuration)
50                    for i in range(len(students_configuration)):
51                        for student in students_copy[i]:
52                            cars_copy[i].add_student(student)
53
54                    if is_everyone_comfy([item for sublist in students_copy for item in sublist]):
55                        for i in range(len(students_configuration)):
56                            for student in students_configuration[i]:
57                                cars_perm[i].add_student(student)
58                        return True
59    except StopIteration as s:
60        return False
61