python 백준 1111 IQ Test

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

로직

• 숫자들 사이의 관계는 nums[i] = a*nums[i-1] + b인 1차 방정식
• N이 1일때는 숫자가 하나 밖에 없으므로 다음 숫자에 어떤 숫자가 나와도 이상하지 않다 => A
• N이 2일때는 두 숫자가 같다면 같은 숫자가 나올 수 밖에 없다
• 하지만 N이 2일때 두 숫자가 다르다면 이 두 수 만으로는 기울기를 구할수 없어 다음에 어떤 숫자가 나오더라도 이상하지 않다 => A
• 숫자들의 관계가 기울기가 a인 1차 방정식 이므로 (nums[i+1] - nums[i]) / (nums[i]-nums[i-1]) 을 이용하여 기울기(a)를 구할 수 있다
• 이렇게 기울기를 구하면 nums[0]과 nums[1]을 이용하여 b값도 구할 수 있다
• 하지만 이러한 기울기가 적용되지 않는 다면 문제 조건만으로는 다음 수를 구할 수 없다 => 'B'

def solution(N:int,nums:list[int])->str or int:
    if N == 1:
        return 'A'
    if N == 2:
        if nums[0] != nums[1]:
            return 'A'
        else:
            return nums[0]
    if nums[0] == nums[1]: # zeroDivision Err 가 일어 날 수 있으므로 0
        a = 0
    else:
        a = int((nums[2] - nums[1]) / (nums[1] - nums[0]))
    b = nums[1] - a * nums[0]
	
    for i in range(1,N):
        if nums[i] != a * nums[i-1] + b:
            return 'B'
    return a*nums[-1] + b

N = int(input())
nums = list(map(int,input().split()))
ans = solution(N,nums)
print(ans)