TCO17 MM R3 PoisonedWine

py.py

import random

class PoisonedWine:
    def testWine(self, num_bottles, num_strips, test_rounds, num_poison):
        max_poison = num_poison
        good_bottles = set()
        bottles = range(num_bottles)

        while test_rounds >= 1 and num_strips:
            r = 1.0 - float(num_poison) / len(bottles)
            max_batch_size = max(1, len(bottles) // num_strips)
            batch_size = max([(s * r ** s, s)
                              for s in range(1, max_batch_size + 1)])[1]

            ratio = 1 if batch_size == max_batch_size else test_rounds * 0.5 + 1
            batch_size = max(int(batch_size / ratio), 1)

            batches = []
            for i in range(num_strips):
                batch = bottles[i * batch_size : (i + 1) * batch_size]
                if batch:
                    batches.append(batch)
            bottles = bottles[batch_size * num_strips:]

            tests = [','.join(map(str, batch)) for batch in batches]
            results = PoisonTest.useTestStrips(tests)

            bad_bottles = []
            for batch, result in zip(batches, results):
                if result == 1:
                    num_strips -= 1
                    num_poison -= 1
                    bad_bottles.extend(batch)
                else:
                    good_bottles.update(batch)

            if num_poison == 0:
                max_bottle = batches[-1][-1]
                while bottles and bottles[0] > max_bottle:
                    good_bottles.add(bottles.pop(0))
                num_poison = max_poison

            random.shuffle(bad_bottles)
            bottles.extend(bad_bottles)
            test_rounds -= 1
        return sorted(set(range(num_bottles)) - good_bottles)