Advent of Code - Day 24a in Python

day24a.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 % 3 == 0
    side_weight = total_weight // 3
    for i in range(1, len(packages) // 3):
        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) // 2):
                if any(get_groups_of_total(packages_2, i_2, side_weight)):
                    found_partition = True
                    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()