počet testů [divide et impera]

#!/bin/env python3.8

import math
import fractions

def n_over_k(n, k):
  assert n >= k
  diff = n - k + 1
  gap_start, gap_end = (k + 1, diff) if k < diff else (diff, k + 1)
  return math.prod(range(gap_end, n + 1)) // math.prod(range(2, gap_start))

class NumberOfTests:
  def __init__(self, depth, cache=None):
    self.cache = cache or {
      (0, 1) : 1,
      # (1, 1) : fractions.Fraction(5, 2),  # unmixed sample subtree hint
    }
    self.depth = depth

  def averageTests(self, infected):
    if not infected:
      return 1
    if (self.depth, infected) in self.cache:
      return self.cache[(self.depth, infected)]
    patients = 2 ** self.depth
    assert infected <= patients
    half_patients = patients // 2
    half_infected = (infected + 1) // 2
    min_infected = max(0, infected - half_patients)
    tests = 0
    weight_sum = 0
    deeper_self = NumberOfTests(self.depth - 1, self.cache)
    for left_infected in range(min_infected, half_infected):
      right_infected = infected - left_infected
      weight = (2 * n_over_k(half_patients, left_infected)
                  * n_over_k(half_patients, right_infected))
      weight_sum += weight
      tests += weight * (
        deeper_self.averageTests(left_infected) +
        deeper_self.averageTests(right_infected))
    if infected % 2 == 0:  # Pascal triangle center
      weight = n_over_k(half_patients, half_infected) ** 2
      weight_sum += weight
      tests += 2 * weight * deeper_self.averageTests(half_infected)
    tests *= fractions.Fraction(1, weight_sum)  # weighted average
    tests += 1  # this vertex
    self.cache[(self.depth, infected)] = tests
    return tests

DEPTH = 6
PATIENTS = 2 ** DEPTH

nt = NumberOfTests(DEPTH)
for i in range(PATIENTS + 1):
  blah = nt.averageTests(i)
  print(f'{i:02}: {float(blah)} ({str(blah)})')