egg_check.py

from typing import Iterable, List


def max_floor(eggs: int, tries: int) -> int:
    if eggs == 1:
        return tries
    elif eggs == 2:
        return tries * (tries + 1) // 2
    else:
        return sum(max_floor(eggs-1, tries-i) for i in range(1, tries)) + tries


def throw_at(eggs_left: int, tries_left: int, floor_range: Iterable[int]) -> List[int]:
    floor_range = list(floor_range)
    if len(floor_range) <= 0:
        return []
    if tries_left == 1:
        return floor_range[:1]
    if eggs_left == 1:
        if len(floor_range) > tries_left:
            print(f"Cannot cover floors {floor_range[tries_left:]}")
        elif len(floor_range) < tries_left:
            print(f"Wasted {tries_left-floor_range} tries")
        return floor_range[:tries_left]
    floors_covered = 0
    throw_at_floors = []
    for i in range(1, tries_left):
        floors_covered += max_floor(eggs_left-1, tries_left-i)+1
        throw_at_floors.append(floors_covered+floor_range[0]-1)
    return throw_at_floors


def main() -> None:
    eggs = int(input("# of eggs : "))
    tries = int(input("# of tries: "))
    assert eggs > 0 and tries > 0
    min_floor_num = 1
    max_floor_num = max_floor(eggs, tries)+1
    floor = 1
    choices = throw_at(eggs, tries, range(min_floor_num, max_floor_num))
    print(f"This configuration should allow {max_floor_num-1} floors")
    while True:
        print(f"Eggs  left: {eggs}")
        print(f"Tries left: {tries}")
        if eggs < 0 or tries < 0:
            print(f"Error: {eggs=} {tries=}")
        if min_floor_num == max_floor_num:
            print(f"Found max floor: {min_floor_num-1}")
            return
        print(f"Floors not checked: ({min_floor_num}, {max_floor_num-1})")
        if eggs == 0 or tries == 0:
            print(f"End")
            return
        if not choices:
            if tries > 1:
                print(f"Wasted {tries-1} tries")
            floors_left = max_floor_num-floor-1
            if floors_left != tries:
                print(f"{floors_left} floors left, but got {tries} tries")
            choices = [floor+i for i in range(1, tries+1)]
        floor = choices.pop(0)
        tries -= 1
        while True:
            broken = input(
                f"Thrown at floor {floor}. Did it break? (y/n): ").lower()
            if broken == "y":
                eggs -= 1
                max_floor_num = floor
                choices = throw_at(eggs, tries, range(
                    min_floor_num, max_floor_num))
                break
            elif broken == "n":
                min_floor_num = floor+1
                break
            else:
                print(f"Invalid response: {broken}")


main()