Egg drop algorithm simulator.

egg-drop-simulator.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
@author Jay Taylor [@jtaylor]
@date 2016-07-17
@description Egg drop algorithm simulator.
"""

import math
import json

class DropContext(object):
    def __init__(self,
        num_floors,
        num_eggs,
        drops=None,
        breaks=None,
        answer=None,
        err=None,
    ):
        self.num_floors = num_floors
        self.num_eggs = num_eggs
        self.drops = drops if drops is not None else []
        self.breaks = breaks if breaks is not None else []
        self.answer = answer
        self.err = err

    def __repr__(self):
        d = dict(
            self.__dict__.items()
            +
            [(f.__name__, f.__call__()) for f in (self.last_broke, self.rem_eggs, self.last_good_floor)]
        )
        s = json.dumps(d)
        return s

    def last_broke(self):
        if len(self.drops) == 0 or len(self.breaks) == 0:
            return False
        broke = self.drops[-1] == self.breaks[-1]
        return broke

    def last_good_floor(ctx):
        last_good = max([0] + [drop for drop in ctx.drops if drop not in ctx.breaks])
        return last_good

    def rem_eggs(self):
        n = self.num_eggs - len(self.breaks)
        return n

def simplest(ctx):
    if ctx.last_broke():
        ctx.answer = ctx.drops[-1]
        return None
    if len(ctx.drops) == 0:
        return 1
    return ctx.drops[-1] + 1

def two_floors_per(ctx):
    if ctx.answer is not None:
        return None
    if len(ctx.breaks) > 1 or (len(ctx.breaks) > 0 and not ctx.last_broke()):
        ctx.answer = ctx.breaks[-1]
        return None
    if ctx.last_broke():
        return ctx.breaks[-1] - 1
    if len(ctx.drops) == 0:
        return 2
    return ctx.drops[-1] + 2

class Stepper(object):
    def __init__(self):
        self.reset()

    def reset(self):
        self.step = -1
        
    def next_floor(self, ctx):
        if ctx.answer is not None:
            return None
        if self.step == -1:
            step = max(1, int(math.sqrt(ctx.num_floors)))
        if len(ctx.breaks) > 0:
            if ctx.last_broke():
                if len(ctx.breaks) > 1:
                    ctx.answer = ctx.breaks[-1]
                    return None
                self.step = 1
        else:
            self.step += 1
        return ctx.last_good_floor() + self.step

def dropper(algo, threshold, num_eggs=2, num_floors=100):
    assert threshold <= num_floors
    ctx = DropContext(num_floors=num_floors, num_eggs=num_eggs)
    while True:
        floor = algo_invoke(algo, ctx)
        if floor is None or ctx.answer is not None:
            break
        ctx.drops.append(floor)
        if floor >= threshold:
            ctx.breaks.append(floor)
            #print(threshold, 'borken', len(ctx.breaks))
        if len(ctx.breaks) >= num_eggs:
            # Provide final opportunity to give an answer.
            algo_invoke(algo, ctx)
            break
    err = None
    if ctx.answer is None and ctx.rem_eggs() == 0:
        err = 'OoE' # Out-of-Eggs.
    elif ctx.answer is None:
        err = 'NOANSWER'
    elif ctx.answer != threshold:
        err = 'INCORRECT(a:%s,g:%s)' % (threshold, ctx.answer)
    ctx.err = err
    return ctx

def tester(algo):
    if isinstance(algo, object) and hasattr(algo, 'reset') and callable(algo.reset):
        algo.reset()
    results = {}
    for floor in range(1, 101):
    #for floor in range(1, 5):
        ctx = dropper(algo, floor)
        results[floor] = ctx #(drops, rem_eggs, error, floors[0:10])
    num_passed = len([result for result in results.values() if result.err == None])
    avg = sum([len(result.drops) for result in results.values()]) / len(results) if len(results) > 0 else 0
    return (algo_name(algo), '%s/%s' % (num_passed, len(results)), avg)

def algo_invoke(algo, ctx):
    if callable(algo):
        return algo(ctx)
    if isinstance(algo, object) and hasattr(algo, 'next_floor') and callable(algo.next_floor):
        return algo.next_floor(ctx)
    raise Exception('Unable to invoke algorithm=%s' % algo)

def algo_name(algo):
    return algo.__name__ if hasattr(algo, '__name__') else algo.__class__.__name__

if __name__ == '__main__':
    from pprint import pprint

    algorithms = [simplest, two_floors_per, Stepper()]
    for algorithm in algorithms:
        pprint(tester(algorithm))