[Baekjoon 백준 1111번] IQ Test Python(파이썬) 풀이

문제

IQ Test의 문제 중에는 공통된 패턴을 찾는 문제가 있다. 수열이 주어졌을 때, 다음 수를 찾는 문제이다.

예를 들어, 1, 2, 3, 4, 5가 주어졌다. 다음 수는 무엇인가? 당연히 답은 6이다. 약간 더 어려운 문제를 보면, 3, 6, 12, 24, 48이 주어졌을 때, 다음 수는 무엇인가? 역시 답은 96이다.

이제 제일 어려운 문제를 보자.

1, 4, 13, 40이 주어졌을 때, 다음 수는 무엇일까? 답은 121이다. 그 이유는 항상 다음 수는 앞 수*3+1이기 때문이다.

은진이는 위의 3문제를 모두 풀지 못했으므로, 자동으로 풀어주는 프로그램을 작성하기로 했다. 항상 모든 답은 구하는 규칙은 앞 수*a + b이다. 그리고, a와 b는 정수이다.

수 N개가 주어졌을 때, 규칙에 맞는 다음 수를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 N이 주어진다. N은 50보다 작거나 같은 자연수이다. 둘째 줄에는 N개의 수가 주어진다. 이 수는 모두 절댓값이 100보다 작거나 같은 정수이다.

출력

다음 수를 출력한다. 만약 다음 수가 여러 개일 경우에는 A를 출력하고, 다음 수를 구할 수 없는 경우에는 B를 출력한다.

 
예제 입력/출력은 아래 링크를 참조하세요.
https://www.acmicpc.net/problem/1111

---

Python 코드

이 문제의 정답 Python 코드는 아래와 같습니다.
저도 처음에는 조금 틀렸습니다 ㅠ

# 변수 선언, 리스트 형태로 받기
T = int(input())
L = [int(x) for x in input().split()]
# 예외조건(A,정답이2개 이상인 경우) 설정
if T == 1 :
    print("A")
    exit() # 만족하면 자동으로 코드 끝내기
if T == 2 :
    if L[0] == L[1] :
        print(L[0])
    elif L[0] != L[1] :
        print("A")
    exit()
    
# a 구하기, 이때 a의 분모가 0이 될 때를 고려하기
if L[1] == L[0] :
    a = 0
elif L[1] != L[0] :
    a = (L[2]-L[1])/(L[1]-L[0])
    
# a 구하기, a가 정수가 아닌 경우를 고려하기
if int(a) != a :
    print("B")
    exit()
# b 구하기
b = L[1]-L[0]*a
# 구한 a와 b가 주어진 값에서 모두 성립하는지 확인, 아니면 (B) 출력 후 코드 끝내기
for i in range (0, T-1) :
    if L[i+1] != L[i]*a+b :
        print("B")
        exit()
    else :
        continue
# 정답 출력하기
print(int(L[T-1]*a+b))

 

---

코드 분석

1. 변수 선언

T = int(input())
L = [int(x) for x in input().split()]

하던대로 첫번째 줄에는 변수 T, 두번째 줄에는 리스트 형태로 연속해서 받을 수 있게 설정해 줍시다.

2. $a, b$ 구하기

주어진 조건을 수열로 표현할 수 있습니다.
 
모든 자연수 $n$에 대해서, $a_{n+1} = a* a_n + b$ 으로 둘 수 있습니다. 따라서 $n$ 대신에 $n+1$을 대입하면
$$a_{n+1} = a* a_n + b$$
$$a_ {n+2} = a* a_ {n+1} + b$$
를 얻습니다. 두 식을 빼면 $b$가 소거되고
$$a_ {n+2} - a_{n+1} = a* a_ {n+1} + a* a_n$$
으로부터
$$a = \frac{a_ {n+2} - a_{n+1}}{a_ {n+1}-a_{n}} , ( n\geq1 ) $$
를 얻을 수 있습니다.
 
위 식은 모든 자연수 $n$에 대해 성립하니까 그냥 아무 자연수를 대입해도 되는데
List 길이에 제한이 있으므로 일단 가장 작은 정수인 $n=1$을 대입하면 (리스트는 0부터 시작)

a = (L[2]-L[1])/(L[1]-L[0])

 
로 표현할 수 있습니다.
그런데 분모가 0인 경우 $a=0$이 되므로 이를 조건문을 활용해서 미리 나눠줍시다.

if L[1] == L[0] :
    a = 0
elif L[1] != L[0] :
    a = (L[2]-L[1])/(L[1]-L[0])

또한 $a$가 정수가 아닐 수도 있으니 정수형과 a의 값이 같은지를 확인하고 다르면 오류(B)를 출력하는 코드를 뒤에 이어붙입니다.

if int(a) != a :
    print("B")
    exit()

 
$b$ 역시 $a$를 안다면, 식 $a_{n+1} = a* a_n + b$ 에 $n=1$을 대입하고 위 값을 대입하면
$$b = a_1-a_0*a$$
임을 알 수 있습니다. 즉

b = L[1]-L[0]*a

입니다.

3. 예외조건

위에서 구한 것은 리스트 원소가 최소 3개 이상 있어야 한다는 것입니다.
그래서
원소가 1개일 경우 : 너무 많음(A)
원소가 2개일 경우 : 두 개가 같으면 그 수 / 다르면 너무 많음(A) 를 출력하게 코드를 짜면 됩니다.

if T == 1 :
    print("A")
    exit()
if T == 2 :
    if L[0] == L[1] :
        print(L[0])
    elif L[0] != L[1] :
        print("A")
    exit()

이런 식으로 말이죠.

4. 확인 및 출력

위에서 구했던 a, b가 주어진 모든 리스트 원소에 대해 만족하는지 따져보고, 아니라면 (B) 출력.
맞다면 다음 항을 출력하는 코드를 짜면 됩니다.

for i in range (0, T-1) :
    if L[i+1] != L[i]*a+b :
        print("B")
        exit()
    else :
        continue

print(int(L[T-1]*a+b))