백준(Baekjoon) 1111번 : IQ Test - python 풀이

🔴 문제 개요

문제 원문 - 백준 온라인 저지 (Baekjoon Online Judge)

🚀 난이도 : GOLD 3

우리가 일상에서 종종 접하는 수열의 규칙을 파악하여 다음 수를 찾아내는 문제입니다. 이 알고리즘만 해결할 수 있다면 문제적 남자에서 우승하는 건 일도 아니겠네요.
단, 수열의 수는 절댓값이 100이 넘지 않는 정수이며, 첫 번째 수 n1과 두 번째 수 n2는 다음과 같은 관계를 가집니다.

${n2 = a*n1 + b}$

이 때 a와 b는 정수입니다.

제한 사항

• 수열의 길이 N은 1부터 50까지 주어질 수 있습니다.

🟠 Solution

수열의 두 수에 대한 점화식을 보면 직선의 방정식이 떠오릅니다. 그렇다면, 직선의 방정식에서 두 점을 알면 기울기에 해당하는 a와 y절편에 해당하는 b를 쉽게 구해낼 수 있겠네요. 이를 위해서 수열에 수가 3개 이상이 있어야, 두 수와의 관계가 2개 이상 정의될 것입니다.

예를 들어, 수열 [x1, x2, x3]과 같이 3개의 수가 있어야 (x1, x2), (x2, x3) 두 관계가 나타날 수 있는 것입니다. 이렇게 두 점이 주어지면, 직선의 방정식의 공식에 따라서 기울기 a와 y절편 b를 구해낼 수 있습니다.

• 기울기 a = (x3-x2) / (x2-x1)
• y절편 b = x2 - x1 * a

이제 a와 b를 구해냈으니, 수열 전체에 점화식을 적용해서 모두 만족하는지 검사하면 문제를 해결할 수 있겠네요! 만약 중간에 수열의 어떤 수가 점화식을 만족하지 않는다면 "B"를 출력하고, 모두 만족한다면 다음 수를 출력해내면 될 것 같습니다.

def solution() -> Union[int, str]:

	...
    
    # Zero Division Error 방지
    a = 0 if nums[1] == nums[0] else (nums[2] - nums[1]) // (nums[1] - nums[0])
    b = nums[1] - nums[0] * a

    for i in range(N - 1):
        # 중간의 수열의 수가 점화식을 만족하지 않는 경우
        if a * nums[i] + b != nums[i + 1]:
            return "B"
    # 모두 만족하는 경우
    return a * nums[-1] + b

a를 정의할 때, 나눠주는 과정에서 Zero Division Error가 발생할 수 있습니다. 생각해보면 문제 특성 상 수열의 두 수가 같으면 a는 0이 되어야되기 때문에 이런 경우 그냥 a를 0으로 둬도 될 것 같습니다.

이렇게 만으로 문제를 해결할 수 있다면 좋겠지만, N은 1부터 주어질 수 있습니다. 우리는 N이 3 이상이라고 가정하고 코드를 작성했기 때문에, N이 1, 2일 때의 예외처리가 필요합니다.

• N=1인 경우
  다음 수가 어떤 수가 나오게 되던 상관 없습니다. 그렇기 때문에 "A"를 반환합니다.
• N=2인 경우
  두 수가 같다면 그 다음 수도 해당 숫자가 나와야 조건을 성립하게 됩니다. 만약 두 수가 다르다면 여러 경우의 수가 나타날 수 있습니다.
  만약 [5, 10]이 주어지면 (a=1, b=5)인 경우와 (a=2, b=0)인 경우 등이 나타날 수 있겠죠. 이 경우 "A"를 반환합니다.

모든 예외처리를 마친 전체 코드는 아래와 같습니다.

import sys
from typing import Union

input = sys.stdin.readline

N = int(input().rstrip())
nums = list(map(int, input().rstrip().split()))


def solution() -> Union[int, str]:
	# N이 1, 2일 때 예외처리
    if N == 1:
        return "A"
    if N == 2:
        if nums[0] == nums[1]:
            return nums[0]
        else:
            return "A"

    # Zero Division Error 방지
    a = 0 if nums[1] == nums[0] else (nums[2] - nums[1]) // (nums[1] - nums[0])
    b = nums[1] - nums[0] * a

    for i in range(N - 1):
        # 중간의 수열의 수가 점화식을 만족하지 않는 경우
        if a * nums[i] + b != nums[i + 1]:
            return "B"
    # 모두 만족하는 경우
    return a * nums[-1] + b


print(solution())

🟡 후기

어느샌가 문제를 접하게 되면 대체 어떤 알고리즘으로 풀어야되지? 라고 먼저 생각하게 되었는데, 특별한 정형화된 알고리즘 없이 우리가 흔히 알고 있는 수학 공식을 접목시켜 문제를 해결할 수 있었기 때문에 굉장히 신선했다는 느낌이 들었습니다. 문제를 대하는 보다 넓은 시야를 갖추게 해준 좋은 문제라는 생각이 들었습니다.

🙏 문제 접근 방법 및 코드에 대한 피드백과 질문은 환영입니다!