Exact meet-in-the-middle solver for minimizing the magnitude of a complex sum in Python with KDTree

Solve the complex-number subset sum exactly in Python. Meet-in-the-middle with KDTree cuts 2^51 brute force to 2^(n/2). Includes numba, SciPy code results.

Choosing, at each index, one of two complex numbers to make the overall sum as close to zero as possible looks innocent until you try it at scale. With two options per position and 51 positions, the search space is 2^51, which explains why a straight brute force approach becomes impractical very quickly. The good news is that you can reframe the task and apply a meet-in-the-middle strategy to get an exact answer dramatically faster.

Problem setup

You have two Python lists of complex numbers. For each index i, you must pick either the element from the first list or the element from the second list. The objective is to minimize the magnitude of the sum of the 51 chosen numbers.

Baseline brute force that stalls

Here is a direct enumeration that explores all assignments. It is easy to read and correct in spirit, but it quickly hits a wall because it must evaluate every one of the 2^51 combinations.

import random
import itertools
import math

alpha = [complex(random.uniform(-10, 10), random.uniform(-10, 10)) for _ in range(51)]
beta = [complex(random.uniform(-10, 10), random.uniform(-10, 10)) for _ in range(51)]

best_abs = math.inf
best_pick = None
for pick_bits in itertools.product([0, 1], repeat=51):
    selection = [alpha[idx] if bit == 0 else beta[idx] for idx, bit in enumerate(pick_bits)]
    s = sum(selection)
    score = abs(s)
    if score < best_abs:
        best_abs = score
        best_pick = selection

print(best_abs)
print(best_pick)

A small syntactic simplification is possible by iterating over product(*zip(...)), but it does not change the complexity story and still leaves you with the same intractable enumeration.

What’s really going on

There is a helpful way to recast the objective. Imagine you start by selecting everything from the first list; call this sum C. Now define a third list delta as the elementwise difference between the second and the first lists, i.e., list3 = list2 - list1. Flipping any index from the first list to the second is equivalent to adding the corresponding delta element to the running sum. The task then becomes to choose a subset of delta that brings the total as close as possible to −C. This matches the Subset Sum Problem.

Because the search is still exponential, the trick is to cut the problem in half and meet in the middle. You split the indices into two halves, enumerate all partial sums in each half, and then pair the two halves so that their sum lands nearest to zero. A KDTree over the complex plane, treated as 2D points, is a convenient way to perform the nearest neighbor lookups for this pairing.

Meet-in-the-middle: from idea to code

The following demonstration shows the meet-in-the-middle structure. It splits the decision set in two, enumerates all sums for each half, and uses KDTree to find the closest-to-zero combined sum. This formulation helps visualize the flow even though it is not tuned for large inputs.

import itertools
import numpy as np
from scipy.spatial import KDTree

def mitm_demo(a_list, b_list):
    n = len(a_list)
    mid = n // 2
    a_head, a_tail = a_list[:mid], a_list[mid:]
    b_head, b_tail = b_list[:mid], b_list[mid:]

    head_accums = [(sum(choice), choice) for choice in itertools.product(*zip(a_head, b_head))]
    tail_accums = [(sum(choice), choice) for choice in itertools.product(*zip(a_tail, b_tail))]

    def to_xy(z):
        return z.real, z.imag

    head_xy = np.array([to_xy(total) for total, _ in head_accums])
    tail_xy = np.array([to_xy(total) for total, _ in tail_accums])

    tree = KDTree(head_xy)
    dists, head_idx_for_tail = tree.query(-tail_xy)

    tail_idx = dists.argmin()
    head_idx = head_idx_for_tail[tail_idx]

    best = head_accums[head_idx][1] + tail_accums[tail_idx][1]
    return abs(sum(best)), list(best)

This cuts the dimensionality and replaces the global pairing with a nearest neighbor search in two dimensions. It is already a big step up from the naive search, but there is more we can do to make it practical for 51 elements.

An optimized implementation that scales to n=51

The optimized version keeps only sums for each half instead of storing the combinations, encodes complex numbers as 2D float arrays for the KDTree, leverages a numba-compiled routine to generate all half-sums, and pulls the exact combination back via more_itertools.nth_product. The approach remains exact. On a typical machine, it solved n=51 in less than 30 seconds.

import itertools
import math
import random
import time

import more_itertools
import numpy as np
from numba import jit
from scipy.spatial import KDTree


def make_complex_list(n):
    return [complex(random.uniform(-10, 10), random.uniform(-10, 10)) for _ in range(n)]


def brute_force_baseline(a_list, b_list):
    assert len(a_list) == len(b_list)
    n = len(a_list)

    best_val = math.inf
    best_sel = None
    for bits in itertools.product([0, 1], repeat=n):
        chosen = [a_list[i] if bit == 0 else b_list[i] for i, bit in enumerate(bits)]
        total = sum(chosen)
        mag = abs(total)
        if mag < best_val:
            best_val = mag
            best_sel = chosen
    return best_val, best_sel


def mitm_demo(a_list, b_list):
    n = len(a_list)
    mid = n // 2
    a_head, a_tail = a_list[:mid], a_list[mid:]
    b_head, b_tail = b_list[:mid], b_list[mid:]

    head_accums = [(sum(choice), choice) for choice in itertools.product(*zip(a_head, b_head))]
    tail_accums = [(sum(choice), choice) for choice in itertools.product(*zip(a_tail, b_tail))]

    def to_xy(z):
        return z.real, z.imag

    head_xy = np.array([to_xy(total) for total, _ in head_accums])
    tail_xy = np.array([to_xy(total) for total, _ in tail_accums])
    tree = KDTree(head_xy)
    dists, head_idx_for_tail = tree.query(-tail_xy)

    tail_idx = dists.argmin()
    head_idx = head_idx_for_tail[tail_idx]
    best = head_accums[head_idx][1] + tail_accums[tail_idx][1]
    return abs(sum(best)), list(best)


@jit(cache=True)
def choice_sum_table(vec_a, vec_b):
    """Generate sums of all pairwise choices from two arrays elementwise.

    Equivalent to: `[sum(current) for current in itertools.product(*zip(vec_a, vec_b))]`
    """
    n = len(vec_a)
    total = 1 << n
    out = np.empty(total, dtype=vec_a.dtype)
    for i in range(total):
        agg = 0.0
        for j in range(n):
            if (i >> (n - j - 1)) & 1:
                agg += vec_b[j]
            else:
                agg += vec_a[j]
        out[i] = agg
    return out


def mitm_accel(a_list, b_list):
    n = len(a_list)
    mid = n // 2
    a_head, a_tail = a_list[:mid], a_list[mid:]
    b_head, b_tail = b_list[:mid], b_list[mid:]

    head_pts = choice_sum_table(np.array(a_head), np.array(b_head))
    tail_pts = choice_sum_table(np.array(a_tail), np.array(b_tail))

    def to_xy_array(arr):
        return arr.view(f"f{arr.real.itemsize}").reshape(-1, 2)

    head_tree = KDTree(to_xy_array(head_pts), balanced_tree=False)
    dists, head_idx = head_tree.query(-to_xy_array(tail_pts), workers=-1)

    tail_idx = dists.argmin()
    head_pick = head_idx[tail_idx]

    head_choice = more_itertools.nth_product(head_pick, *zip(a_head, b_head))
    tail_choice = more_itertools.nth_product(tail_idx, *zip(a_tail, b_tail))
    best = head_choice + tail_choice
    return abs(sum(best)), list(best)


def self_test():
    for n in range(2, 15):
        for seed in range(10):
            random.seed(seed)
            a = make_complex_list(n)
            b = make_complex_list(n)

            exp_val = brute_force_baseline(a, b)
            got_demo = mitm_demo(a, b)
            assert exp_val == got_demo, f"Mismatch in demo MITM: n={n}, seed={seed}"

            got_fast = mitm_accel(a, b)
            assert exp_val == got_fast, f"Mismatch in optimized MITM: n={n}, seed={seed}"
    print("All tests passed!")


def perf():
    n = 51
    random.seed(0)
    a = make_complex_list(n)
    b = make_complex_list(n)
    t0 = time.perf_counter()
    _ = mitm_accel(a, b)
    dt = time.perf_counter() - t0
    print(f"n={n}, elapsed={dt:.0f} sec")


if __name__ == "__main__":
    self_test()
    perf()

This exact approach reframes the selection as a subset sum, enumerates both halves independently, and finds the nearest opposing sums with KDTree. It keeps the heavy lifting bounded to 2^(n/2) per half rather than 2^n overall, and then stitches the best pair back to the full assignment. The code validates correctness against the full brute-force baseline for small n and demonstrates that n=51 completes in under half a minute on a regular PC.

Why this matters

You have only two knobs at each position, but that is more than enough to explode the search space. Direct enumeration is correct yet non-starter for n=51 because it must consider all 2^51 variants. The meet-in-the-middle strategy offers the same exactness with a dramatically more manageable inner loop, and using KDTree for nearest neighbor search over the complex plane keeps the pairing step simple. Small engineering choices like not materializing combinations, leveraging numba for the sum table, and reconstructing the winning assignment via more_itertools.nth_product are what make the method practical end-to-end.

If the selection must be made across more than two lists at once, the situation fundamentally changes. With ten lists of 51 elements, the task becomes impossible to solve as posed, and the right move is to pursue a way to reach your broader goal without solving this exact problem.

There are also alternative directions for those who do not require an exact optimum. One can approximate by drawing random assignments, flipping indices when that decreases the magnitude, and stopping when no single flip helps, tracking the best magnitude seen along the way. Another angle is to formulate it as MIQP, which opens the door to solvers that can tackle mixed-integer quadratic programs.

Conclusion

When faced with choosing between two complex values per index and minimizing the norm of the total, the temptation to brute force is understandable but unworkable for 51 decisions. Reframing the task as a subset sum and applying meet-in-the-middle with a KDTree turns it into an exact and tractable solver. For the two-list case, this approach is both principled and fast in practice. If your requirements grow beyond two lists, seek a different path altogether or accept an approximation strategy. Whatever you choose, begin by measuring where the time goes and align the technique with your tolerance for compute and exactness.