[Python] 백준 1111번: IQ Test

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

1111번: IQ Test

 

 > 접근

초기에는 N <= 2를 기준으로 나누어, 참일 경우에는 무조건 A를 출력하고, 아닌 경우에는 a, b를 구한 뒤 검증하여 다음 수를 출력하거나 B를 출력하는 정도로 하였다.

하지만 예외가 좀 있었다. 첫번째로 N <= 2가 참일 때의 내부 로직에 대한 것이다. [57, 57]이 들어왔을 때인데, 이 경우에는 57이 다음 수로 나와야하기 때문에, (N == 2 and nums[0] == nums[1])라는 조건문을 두어 해결하였다.

다음으로 N <= 2가 거짓일 때의 예외이다. [1, 1, 2, 3, 4] 또는 [1, 1, 1, 1, 1]와 같은 경우가 그러하다. 이 때는 a, b의 값을 0, nums[0]로 넣어주면 된다. [1, 1, 1, 1, 1]에서 a와 b에 각각 0, 1을 넣었을 때, prev*0 + 1 = next 이므로 next = 1이 성립하기 때문이다.

 

# -*- coding: utf-8 -*-
import sys
input = sys.stdin.readline

def is_a_and_b_valid(a, b):
  for i in range(N-1):
    if(nums[i]*a + b != nums[i+1]):
      return False

  return True

if __name__ == '__main__':
  N = int(input())
  nums = list(map(int, input().split()))
  if(N <= 2):
    if(N == 2 and nums[0] == nums[1]):
      # ex. [57, 57]
      print(nums[0])
    else:
      # 일반적인 경우
      print("A")
  else:                             
    if(nums[0] == nums[1]):    
      # ex. [57, 57, ...]
      a = 0
      b = nums[0]
    else:
      # 일반적인 경우
      # nums[0]*a + b = nums[1]
      # nums[1]*a + b = nums[2]
      # 따라서, (nums[1] - nums[0])*a = nums[2] - nums[1]이고, 이를 전개하면
      # a = (nums[2] - nums[1]) / (nums[1] - nums[0])이다.
      a = int((nums[2] - nums[1]) / (nums[1] - nums[0]))
      b = nums[1] - nums[0]*a    
      
    if(is_a_and_b_valid(a, b)):
      print(nums[-1]*a + b)
    else:
      print("B")