Maximum combination health of AI modules

A combination's health is sum(health) * min(stability). Two forces pull
against each other: adding a module raises the sum but can only lower the
minimum stability. The trick is to stop treating the minimum as something
the combination produces and start treating it as something we pick.

Brute force. Try every subset of size 1..k, compute sum and minimum, keep
the best. That is O(2^n * n) and is only useful as a checker for n <= 15.
Say it out loud, then move on.

Fix the bottleneck. Suppose module j is the least stable module in the
chosen combination. Then every other chosen module has stability >= s[j],
and the score is s[j] * (h[j] + the other chosen healths). To maximize it
we want the up to k - 1 largest healths among modules at least as stable
as j, and we must include j itself.

So sort modules by stability, descending, and sweep. When we reach module
j, every module already seen is at least as stable. Keep a min-heap of
the healths of the best modules seen so far, capped at k, along with
their running sum:

1. Push h[j] and add it to the sum.
2. If the heap now holds more than k, pop the smallest health and
   subtract it.
3. Score the candidate sum * s[j].

Why it is correct even when step 2 pops h[j] itself. Then the heap
holds k modules that are all at least as stable as j, and their real
minimum stability is >= s[j]. The candidate sum * s[j] therefore
underestimates a combination that really exists, so it never overstates
the answer. And the true optimum is scored exactly at its own bottleneck
module j*: at that step the heap holds the k largest healths among
modules with stability >= s[j*], which is at least as good as the optimal
set.

Invariants. After processing a prefix of the sorted order, the heap
holds the min(k, prefix length) largest healths in that prefix and
total equals their sum. All healths are non-negative, so a larger sum is
never worse.

Ties in stability need no special handling: any order among equal
stabilities gives the same candidates.

import heapq
from itertools import combinations
from typing import List


def max_combination_health(health: List[int], stability: List[int], k: int) -> int:
    """Best value of sum(health) * min(stability) over 1..k distinct modules."""
    n = len(health)
    if n != len(stability):
        raise ValueError("health and stability must have the same length")
    if n == 0 or k <= 0:
        return 0
    k = min(k, n)

    # Most stable first: each module in turn is the bottleneck candidate.
    order = sorted(range(n), key=lambda i: stability[i], reverse=True)

    heap: List[int] = []  # healths of the chosen modules, smallest on top
    total = 0
    best = 0
    for i in order:
        heapq.heappush(heap, health[i])
        total += health[i]
        if len(heap) > k:
            total -= heapq.heappop(heap)
        best = max(best, total * stability[i])
    return best


def brute_force(health: List[int], stability: List[int], k: int) -> int:
    best = 0
    n = len(health)
    for size in range(1, min(k, n) + 1):
        for combo in combinations(range(n), size):
            s = sum(health[i] for i in combo)
            m = min(stability[i] for i in combo)
            best = max(best, s * m)
    return best


if __name__ == "__main__":
    import random

    assert max_combination_health([4, 9, 2, 6], [5, 2, 8, 4], 2) == 40
    assert max_combination_health([4, 9, 2, 6], [5, 2, 8, 4], 3) == 48
    assert max_combination_health([7], [3], 1) == 21
    assert max_combination_health([], [], 3) == 0
    assert max_combination_health([5, 5], [1, 1], 0) == 0
    assert max_combination_health([1, 2, 3], [10, 10, 10], 5) == 60
    assert max_combination_health([0, 0], [9, 9], 2) == 0

    rng = random.Random(7)
    for _ in range(2000):
        n = rng.randint(1, 8)
        h = [rng.randint(0, 20) for _ in range(n)]
        s = [rng.randint(0, 20) for _ in range(n)]
        k = rng.randint(1, n + 1)
        assert max_combination_health(h, s, k) == brute_force(h, s, k)
    print("ok")

• Sorting: O(n log n).
• Sweep: n pushes and at most n pops on a heap of size k + 1, so
  O(n log k).
• Total: O(n log n) time, O(n) space for the sorted order plus O(k) for
  the heap.
• The brute-force checker is O(2^n * n) and is for testing only.

• health=[4, 9, 2, 6], stability=[5, 2, 8, 4], k=2 -> 40
  (modules 0 and 3: (4 + 6) * 4; module 1 has the most health but its
  stability of 2 drags any combination down).
• Same input, k=3 -> 48 (modules 0, 2, 3: (4 + 2 + 6) * 4).
• One module, health=[7], stability=[3], k=1 -> 21.
• Empty input -> 0.
• k=0 -> 0 (no combination allowed).
• k larger than n, equal stabilities: health=[1, 2, 3],
  stability=[10, 10, 10], k=5 -> 60 (take all).
• All healths zero -> 0.
• A single very stable module can beat a large unstable group:
  health=[10, 10, 10], stability=[1, 1, 100], k=3 -> 1000.
• Randomized: 2,000 random cases with n <= 8 agree with the brute force.

1. Exactly k modules, not up to k. Only score a candidate once the
   heap holds exactly k items. Return -1 (or raise) if n < k.
2. Large values; answer modulo 1e9+7. Compare the true products (Python
   ints are unbounded; in Java or C++ use 64-bit or 128-bit for the
   comparison) and apply the modulo only to the final answer. Taking the
   modulo before comparing gives wrong maxima.
3. Return the chosen modules, not just the score. Store (health, index) in the heap and remember the step that produced best. Rebuild
   the set by re-running the sweep up to that step, or copy the heap when
   best improves (O(k) per improvement).
4. Modules arrive as a stream and queries ask for the current best.
   The sorted-by-stability sweep no longer fits. Keep modules in a
   structure ordered by stability (a balanced BST or a segment tree over
   compressed stability values) that can answer "sum of the top k - 1
   healths among modules with stability >= s". Each query then scans the
   candidate bottlenecks, which costs more; discuss rebuilding in batches
   if writes are rare.
5. Combination score uses the sum of stability and the min of health
   instead. The structure is symmetric: swap which attribute you sort by
   and which one the heap tracks.

Mid-level (L4)

• States the brute force and its cost before optimizing.
• Finds the "fix the minimum" idea with a hint, and codes the sort plus
  heap sweep correctly.
• Handles k > n, empty input and k = 0 without prompting after a
  test walk-through.

Senior (L5)

• Reaches the bottleneck sweep on their own and explains why popping the
  current module still gives a safe lower bound.
• States the invariant on the heap and its sum, and gives O(n log n) time.
• Tests with a randomized brute-force comparison or reasons through the
  tie and zero cases unprompted.

Staff (L6+)

• Connects the problem to the general pattern (maximize sum times min by
  enumerating the min) and adapts it quickly to the follow-ups.
• Discusses overflow and modulo pitfalls and how they differ by language.
• Sketches a sound design for the streaming version and its trade-offs
  between query and update cost.