간단한 숫자야구 구현

number_baseball.py

import itertools
import random


class NumberBaseball(object):
    RANGE = range(10)
    PITCHES = 4

    def __init__(self, ans):
        assert len(ans) == len(set(ans))
        assert all(isinstance(x, int) for x in ans)
        self.ans = ans

    def defence(self, comb, ans=None):
        if ans is None:
            ans = self.ans
        strikes, balls = 0, 0
        assert len(ans) == len(comb)
        for i, x in enumerate(comb):
            if x == ans[i]:
                strikes += 1
            elif x in ans:
                balls += 1
        return strikes, balls

    def _filter(self, candidates, comb, defence_result):
        for cand in candidates:
            if self.defence(comb, cand) == defence_result:
                yield cand

    def guess(self):
        comb = random.sample(self.RANGE, self.PITCHES)
        candidates = itertools.permutations(self.RANGE, self.PITCHES)
        while True:
            res = self.defence(comb)
            print(f"{tuple(comb)} - {res[0]}S {res[1]}B")
            candidates = list(self._filter(candidates, comb, res))
            if len(candidates) == 1:
                return candidates[0]
            elif len(candidates) == 0:
                raise Exception("No possible answer")
            comb = random.choice(candidates)
        
            
if __name__ == "__main__":
    ans = random.sample(NumberBaseball.RANGE, NumberBaseball.PITCHES)
    print(f"Generated answer: {tuple(ans)}")
    game = NumberBaseball(ans)
    print(f"Guessed answer: {game.guess()}")