[백준/1111] IQ Test (Python)

Problem : https://www.acmicpc.net/problem/1111

1111번: IQ Test

 

Difficulty : Gold 3

 

Status : Solved

 

Time : 00:29:38

 

---

 

문제 설명

 

 

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를 출력한다.

 

입력 예시

 

4
1 4 13 40

 

출력 예시

 

121

 

 

---

 

풀이

 

생각보다 오래 걸렸던 문제. 예외처리를 조금 신경써줄 필요가 있다.

 

N = 1이라면, 이 뒤에 수많은 수가 올 수 있으므로 "A"를 출력한다.

 

N = 2이라면 조건을 만족하는 (a, b)쌍 자체는 무한하다. 따라서 보통의 경우는 "A"를 출력한다. 단, 첫번째 수와 두 번째 수가 같다면 (a, b)쌍이 어떻던 간에 다음 숫자는 앞의 두 숫자와 동일하다.

 

N >= 3인 경우는 연립방정식을 통해 풀이할 수 있다.

• i번째 수를 N[i]라고 둔다면, a * N[0].+ b = N[1], a * N[1] + b = N[2] 가 성립한다.
• 이를 풀면, a = (N[2] - N[1]) / (N[1] - N[0]) 이 성립한다.
  • 첫째로, N[1] == N[0]인 특수한 경우가 있을 수 있다. 이 경우는 N = 2인 경우와 동일하다. 만약 세 번째 수 N[2] 역시 같다면 다음 숫자는 앞의 두 숫자와 동일하다는 전제를 깔 수 있으므로, 무한한 a, b쌍 중 (1, 0)과 같은 값을 정한다. 아니라면 "B"를 출력한다.
  • 둘째로, N[1] != N[0]인 경우는 a의 유일성이 보장된다. 또한 a는 정수라는 조건이 있으므로, 만약 정수가 아니라면 "B"를 출력한다. b는 a에 따라 자동으로 유도된다.
• 앞서 구한 (a, b)쌍이 모든 N에 대해 성립하는지 검사한다. 만약 검사시 결과가 맞지 않다면, "B"를 출력한다.
• 모든 검사를 통과하였다면 다음 예측값인 a * N[-1] + b를 출력하도록 한다.

 

 

풀이 코드

N = int(input())
num_list = list(map(int, input().split()))
if N == 1 :
  print("A")
elif N == 2 :
  print("A" if num_list[0] != num_list[1] else num_list[0])
else :
  if num_list[1] == num_list[0] :
    if num_list[2] != num_list[1] :
      print("B")
      exit()
    a, b = 1, 0
  else :
    a = (num_list[2] - num_list[1]) / (num_list[1] - num_list[0])
    if not a.is_integer() :
      print("B")
      exit()
    a = int(a)
    b = num_list[1] - num_list[0] * a
  for i in range(1, N) :
    if num_list[i-1] * a + b != num_list[i] :
      print("B")
      exit()
  print(num_list[-1] * a + b)

풀이 완료!