covid_bs_algorithm.py

import math

DEBUG = False


def generate_input_sample(params):
    population = COVIDPopulation(params["N"], params["K"])

    return population


def covid_binary_search(population, start, end, costly):
    if start > end: return

    # Test this whole population segment
    # Practically speaking, mix blood and do the COVID test
    if DEBUG: print(f"Testing: {start} to {end}")
    costly()
    if not population.is_infected(start, end):
        return

    # No more splitting if start == end. We just tested an individual.
    if start == end: return

    # Split into two segments [start, mid] and [mid+1, end]
    mid = math.floor((start+end)/2)
    covid_binary_search(population, start, mid, costly)
    covid_binary_search(population, mid+1, end, costly)


def covid_binary_search_wrapper(input_sample, costly):
    population = input_sample
    n = population.N

    return covid_binary_search(population, 0, n-1, costly)