[프로그래머스] 숫자 야구

문제 요약

숫자 야구는 서로 다른 숫자 4개로 이루어진 비밀번호를 맞히는 게임이다.

비밀번호는 다음 조건을 만족한다.

• 숫자는 1부터 9까지 사용한다.
• 같은 숫자는 중복해서 사용할 수 없다.
• 총 4자리 숫자다.

예를 들어 가능한 비밀번호는 다음과 같다.

1234
1357
9876

하지만 다음은 불가능하다.

1123  # 1이 중복됨
0234  # 0이 포함됨

우리는 submit() 함수를 호출해 숫자를 제출할 수 있다.

제출한 숫자에 대해 다음과 같은 단서가 주어진다.

• 숫자와 위치가 모두 같으면 STRIKE
• 숫자는 포함되어 있지만 위치가 다르면 BALL
• 숫자가 포함되어 있지 않으면 OUT

단서는 다음 형식으로 반환된다.

"xS yB"

예를 들어 "2S 1B"는 스트라이크 2개, 볼 1개라는 의미다.

목표는 submit() 호출 횟수가 n번을 넘지 않도록 비밀번호를 찾아 반환하는 것이다.

핵심 아이디어

가능한 비밀번호 후보는 많지 않다.

비밀번호는 1~9 중 서로 다른 숫자 4개로 이루어져 있다.

따라서 가능한 후보 수는 다음과 같다.

9P4 = 9 x 8 x 7 x 6 = 3024

3024개는 충분히 작기 때문에 모든 후보를 미리 만들어둘 수 있다.

그다음 submit()으로 단서를 받을 때마다, 그 단서와 모순되는 후보를 제거한다.

즉, 풀이 흐름은 다음과 같다.

1. 가능한 모든 비밀번호 후보를 만든다.
2. 후보 하나를 골라 submit()에 제출한다.
3. 반환된 단서를 파싱한다.
4. 해당 단서와 일치하지 않는 후보를 제거한다.
5. 정답을 찾을 때까지 반복한다.

가능한 후보 만들기

파이썬의 itertools.permutations를 사용하면 서로 다른 숫자 4개로 이루어진 순열을 쉽게 만들 수 있다.

from itertools import permutations

candidates = [
    "".join(number)
    for number in permutations("123456789", 4)
]

예를 들어 생성되는 후보는 다음과 같다.

1234
1235
1236
...
9876

각 후보는 문자열로 관리한다.

문자열로 관리하면 자릿수 비교가 쉽기 때문이다.

단서 계산 함수

후보를 제거하려면 두 숫자를 비교했을 때 몇 스트라이크, 몇 볼이 나오는지 직접 계산할 수 있어야 한다.

def get_result(guess, answer):
    strike = 0
    ball = 0

    for i in range(4):
        if guess[i] == answer[i]:
            strike += 1
        elif guess[i] in answer:
            ball += 1

    return strike, ball

guess는 제출한 숫자이고, answer는 정답 후보라고 생각하면 된다.

각 자릿수를 비교하면서 다음을 확인한다.

if guess[i] == answer[i]:

숫자와 위치가 모두 같으면 스트라이크다.

elif guess[i] in answer:

숫자는 포함되어 있지만 위치가 다르면 볼이다.

단서 계산 예시

정답 후보가 "1357"이고 제출값이 "3457"이라고 하자.

각 자리 비교는 다음과 같다.

위치	제출 숫자	정답 숫자	판정
0	3	1	BALL
1	4	3	OUT
2	5	5	STRIKE
3	7	7	STRIKE

따라서 결과는 다음과 같다.

(2, 1)

이는 2S 1B를 의미한다.

submit 결과 파싱하기

submit() 함수는 문자열을 반환한다.

예를 들어 다음과 같은 형태다.

"2S 1B"

이를 숫자로 바꾸기 위해 파싱 함수를 만든다.

def parse_result(result):
    strike_text, ball_text = result.split()
    strike = int(strike_text[:-1])
    ball = int(ball_text[:-1])
    return strike, ball

"2S"에서 마지막 문자 S를 제외하면 "2"가 된다.

strike_text[:-1]

"1B"도 같은 방식으로 "1"만 꺼낼 수 있다.

후보 제거하기

어떤 숫자를 제출했을 때 결과가 (strike, ball)이라고 하자.

그렇다면 실제 정답은 제출값과 비교했을 때 반드시 같은 결과를 만들어야 한다.

따라서 후보 중 다음 조건을 만족하는 것만 남긴다.

get_result(guess, candidate) == result

코드로는 다음과 같다.

candidates = [
    candidate
    for candidate in candidates
    if get_result(guess, candidate) == result
]

이 과정을 반복하면 가능한 후보가 점점 줄어든다.

전체 코드

from itertools import permutations


def solution(n, submit):
    def get_result(guess, answer):
        strike = 0
        ball = 0

        for i in range(4):
            if guess[i] == answer[i]:
                strike += 1
            elif guess[i] in answer:
                ball += 1

        return strike, ball

    def parse_result(result):
        strike_text, ball_text = result.split()
        strike = int(strike_text[:-1])
        ball = int(ball_text[:-1])
        return strike, ball

    candidates = [
        "".join(number)
        for number in permutations("123456789", 4)
    ]

    for _ in range(n):
        guess = candidates[0]
        result = parse_result(submit(int(guess)))

        if result == (4, 0):
            return int(guess)

        candidates = [
            candidate
            for candidate in candidates
            if get_result(guess, candidate) == result
        ]

    return int(candidates[0])

코드 설명

1. 후보 생성

candidates = [
    "".join(number)
    for number in permutations("123456789", 4)
]

1~9 중 서로 다른 숫자 4개를 뽑아 만들 수 있는 모든 후보를 생성한다.

후보 수는 3024개다.

2. 후보 하나 제출

guess = candidates[0]
result = parse_result(submit(int(guess)))

현재 남아 있는 후보 중 첫 번째 후보를 제출한다.

후보는 문자열로 관리하지만, submit()에는 정수를 제출해야 하므로 int()로 변환한다.

3. 정답이면 반환

if result == (4, 0):
    return int(guess)

4자리 숫자가 모두 위치까지 맞으면 4S 0B다.

이 경우 비밀번호를 찾은 것이므로 바로 반환한다.

4. 모순되는 후보 제거

candidates = [
    candidate
    for candidate in candidates
    if get_result(guess, candidate) == result
]

제출한 숫자와 같은 단서를 만들 수 없는 후보는 정답이 될 수 없다.

따라서 단서와 일치하는 후보만 남긴다.

시간 복잡도

가능한 후보 수는 고정되어 있다.

9P4 = 3024

각 제출마다 모든 후보를 한 번씩 검사한다.

비교할 때는 길이 4짜리 문자열만 확인한다.

따라서 한 번의 후보 필터링 비용은 다음과 같다.

O(3024 x 4)

상수 수준으로 볼 수 있다.

제출 횟수를 n이라고 하면 전체 시간 복잡도는 다음과 같다.

O(n x 3024)

코딩테스트 환경에서는 충분히 빠르다.

공간 복잡도

가능한 후보를 저장한다.

O(3024)

이 역시 상수 수준이다.

주의할 점

0은 사용하지 않는다

문제에서 비밀번호는 1~9 사이의 숫자로만 이루어진다.

따라서 후보 생성 시 "0123456789"가 아니라 "123456789"를 사용해야 한다.

permutations("123456789", 4)

숫자는 중복되지 않는다

비밀번호는 서로 다른 숫자 4개로 구성된다.

따라서 중복을 허용하는 중첩 반복문보다 permutations를 사용하는 것이 깔끔하다.

submit 호출 횟수 제한

submit() 호출 횟수는 n을 넘으면 안 된다.

따라서 반복문은 다음처럼 작성한다.

for _ in range(n):

정리

이 문제는 가능한 비밀번호 후보가 3024개뿐이라는 점을 이용한다.

풀이 핵심은 다음과 같다.

1. 가능한 모든 비밀번호 후보를 만든다.
2. 후보 하나를 제출해 단서를 받는다.
3. 단서와 일치하지 않는 후보를 제거한다.
4. 4S 0B가 나올 때까지 반복한다.

숫자 야구 문제는 단서를 직접 계산할 수 있기 때문에, 후보군을 계속 줄여나가는 방식으로 간단하게 해결할 수 있다.