1111. IQ Test (Python)

1111번: IQ Test

 

풀이

 방정식을 통하여 a와 b를 구하고 각 숫자에 * a + b를 했을 때 다음 숫자가 나오는지를 확인하면 된다.

1, 4, 13, 40

 

 예를 들어 위의 예제에서 1a + b = 4와 4a + b =13의 방정식을 사용하여 a와 b를 구한다. 두 식을 빼면 3a = 9가 나오고 여기에서 a=3이 됨을 알 수 있다. 이를 통해 b는 1이 된다. 이제 앞에서 부터 차례대로 * 3 + 1을 하여 다음 숫자가 나오는지 확인한다.

 

 문제를 풀 때 몇 가지 예외 처리를 해줘야 한다.

 

1. N이 1인 경우
   -> 뒤에 아무 숫자나 올 수 있으므로 'A'를 출력한다.
   
   
2. N이 2인 경우
   1. 첫 번째 숫자와 두 번째 숫자가 같은 경우
      -> 다음 숫자 역시 같은 숫자가 나온다.
      
      
   2. 두 숫자가 다른 경우
      -> a와 b를 무한히 조합하여 만들 수 있다. 그러므로 다음 수는 여러 개일 수 있다.
      
      
3. N이 3 이상이고 첫 번째 숫자 == 두 번째 숫자 == 세 번째 숫자일 경우
   -> 해당 숫자가 몇 개가 있는지 카운트한다. 만약 똑같은 숫자가 N개 존재한다면 다음 수 역시 해당 숫자가 되고 만약 다른 숫자가 존재할 경우 다음 수를 구할 수 없으므로 'B'를 출력한다.
   
   

코드

import sys

# 모든 숫자에 * a + b 연산을 하여 다음 숫자가 나오는지 체크한다.
def solve(a, b):
    for i in range(N-1):
        # 다음 숫자가 나오지 않으면 바로 False를 리턴한다.
        if nums[i] * a + b != nums[i+1]:
            return False 

    return True


if __name__ == '__main__':
    N = int(input())
    nums = list(map(int, input().split()))

    # N이 1일 경우에는 다음 수가 여러 개일 수 있다
    if N == 1:
        print('A')
        sys.exit(0)

    # N이 2이고
    if N == 2:
        if nums[0] == nums[1]: # 두 숫자가 같을 경우 다음 숫자 역시 같은 숫자가 온다.
            print(nums[0])
        else: # 두 숫자가 다르다면 다음 수가 여러 개일 수 있다.
            print('A')
        sys.exit(0)

    # N이 3 이상이고 앞의 세 개의 숫자가 같을 경우
    if nums[0] == nums[1] == nums[2]:
        if nums.count(nums[0]) == N: # 해당 숫자가 N개 존재한다면 다음 수 역시 해당 숫자가 된다
            print(nums[0])
        else: # 그렇지 않으면 다음 수를 구할 수 없다
            print('B')
        sys.exit(0)

    a, b = 0, 0
    # 단순히 앞의 세 숫자를 사용하여 a와 b를 구할 경우 0으로 나누는 경우가 발생할 수 있다.
    # for문을 통해서 a와 b를 구한다.
    for i in range(N-2): 
        if nums[i+1] - nums[i] == 0:
            continue
        else: 
            a = (nums[i+2] - nums[i+1]) // (nums[i+1] - nums[i])
            b = nums[i+1] - (nums[i] * a)
            break
    
    # 모든 숫자가 규칙에 맞으면 해당 규칙에 맞춰 다음 수를 구한다
    # 그렇지 않은 경우에는 다음 수를 구할 수 없으므로 'B'를 출력한다.
    if solve(a, b):
        print(nums[N-1] * a + b)
    else:
        print('B')