search.py

import itertools
import pprint
import typing

_SELECTED = 5

_MAX_VALUE = 40
values = list(range(1, _MAX_VALUE + 1))

_REQUIRED_SUM = _MAX_VALUE // 2


def _brute_force():
    all_elements = filter(
        lambda d: sum(d) == _REQUIRED_SUM,
        itertools.product(*(values for _ in range(_SELECTED)))
    )

    for i in all_elements:
        print(i)


def _graph():
    Node = typing.NamedTuple(
        'Node',
        (
            ('layer', 'int'),
            ('offset', 'int'),
            ('parent', typing.Optional['Node'])
        )
    )
    layers = [values for _ in range(_SELECTED)]

    def _get_children(e: Node) -> typing.List[Node]:
        if e.layer >= _SELECTED - 1:
            return []
        parent_layer, parent_offset = e.layer, e.offset
        return [Node(parent_layer + 1, i, e) for i in range(len(layers[parent_layer + 1]))]

    def _count(e: typing.List[Node]) -> int:
        return sum(
            map(lambda d: layers[d.layer][d.offset] if d.layer >= 0 else 0, e)
        )

    def _solution(e: typing.List[Node]) -> bool:
        return _count(e) == _REQUIRED_SUM and len(e) == _SELECTED

    def _ok(e: typing.List[Node]) -> bool:
        return _count(e) < _REQUIRED_SUM <= _MAX_VALUE * (_SELECTED - len(e))

    def _print_elements(e: typing.List[Node]):
        pprint.pprint(
            [layers[d.layer][d.offset] for d in e if d.layer >= 0])

    def _get_path(e: Node) -> typing.List[Node]:
        result = []
        while e.layer >= 0:
            result.append(e)
            e = e.parent
        return result[::-1]

    elements = [Node(-1, -1, None)]
    while elements:
        node = elements.pop()
        path = _get_path(node)
        if _solution(path):
            _print_elements(path)
        elif not _ok(path):
            continue
        children = _get_children(node)
        elements.extend(children)


_graph()