Advent of Code - Day 24b in Python

day24b.py

#!/usr/bin/env python3

from functools import reduce

def get_groups_of_total(collection, amount, total):
    """The collection must be sorted in ascending order"""
    for (pos, elem) in enumerate(collection):
        if elem > total:
            break

        if amount == 1:
            if elem == total:
                yield [elem]
        else:
            reduced_collection = collection[pos+1:]
            for group in get_groups_of_total(reduced_collection, amount - 1, total - elem):
                yield [elem] + group

def calc_qe(group):
    return reduce(lambda x, y: x*y, group, 1)

def main():
    input_fname = 'day24_input.txt'
    packages = []
    with open(input_fname) as f:
        for line in f:
            packages.append(int(line.strip()))

    # Test input
    # packages = [1, 2, 3, 4, 5, 7, 8, 9, 10, 11]

    packages.sort()
    min_size = None
    min_qe = None

    total_weight = sum(packages)
    assert total_weight % 4 == 0
    side_weight = total_weight // 4
    for i in range(1, len(packages) // 4):
        if min_size and i > min_size:
            break
        for group in get_groups_of_total(packages, i, side_weight):
            packages_2 = packages[:]
            for x in group:
                packages_2.remove(x)

            found_partition = False

            for i_2 in range(len(group), len(packages_2) // 3):
                for group_2 in get_groups_of_total(packages_2, i_2, side_weight):
                    packages_3 = packages_2[:]
                    for x in group_2:
                        packages_3.remove(x)
                    for i_3 in range(len(group_2), len(packages_3) // 2):
                        if any(get_groups_of_total(packages_3, i_3, side_weight)):
                            found_partition = True
                            break
                    if found_partition:
                        break
                if found_partition:
                    break

            if found_partition:
                if min_size is None or len(group) < min_size:
                    min_size = len(group)
                    qe = calc_qe(group)
                    print('Found partition with min size {} and QE {}'.format(min_size, qe))
                    if min_qe is None or qe < min_qe:
                        min_qe = qe

    print('Min QE:', min_qe)

if __name__ == '__main__':
    main()