Search code examples
pythonalgorithmlanguage-agnostic

Bitmask with weights


I have following pairs:

pairs = [(1,9),(5,5),(6,4),(2,8)]

From the pair I can take only one element. I need to generate n lists of binary mask combinations with minimal sum, where 0 for first element and 1 for second.

Output of provided example if n == 5:
[0,0,1,0], [0,1,1,0], [0,0,0,0], [0,1,0,0], [0,0,1,1]

Is it kind of Knapsack Problem? What is the best algorithm for doing this?


Solution

  • The idea behind the code below is to enumerate the sorted combinations recursively (well, not using recursion) for the tail of the pairs list, then do a sorted merge between the combinations that use the lesser element of the first pair and the combinations that use the greater element of the first pair. The space requirement is minimized using generators (should generally be on the order of the number of elements you take).

    import collections
    import itertools
    
    
    def extend_sorted(pair, sequence):
        argmin_pair = min(range(2), key=pair.__getitem__)
        min_pair = pair[argmin_pair]
        argmax_pair = 1 - argmin_pair
        max_pair = pair[argmax_pair]
        buffer = collections.deque()
        for sub_combo, sub_total in sequence:
            while buffer and buffer[0][1] < min_pair + sub_total:
                yield buffer.popleft()
            yield ([argmin_pair] + sub_combo, min_pair + sub_total)
            buffer.append(([argmax_pair] + sub_combo, max_pair + sub_total))
        while buffer:
            yield buffer.popleft()
    
    
    def enumerate_least_sums(pairs):
        sequence = [([], 0)]
        for pair in reversed(pairs):
            sequence = extend_sorted(pair, sequence)
        for combo, total in sequence:
            yield combo
    
    
    if __name__ == "__main__":
        print(*itertools.islice(enumerate_least_sums([(1, 9), (5, 5), (6, 4), (2, 8)]), 5))
    

    Output:

    [0, 0, 1, 0] [0, 1, 1, 0] [0, 0, 0, 0] [0, 1, 0, 0] [0, 0, 1, 1]