[백준] 1111번: IQ Test - [Python]

 

 

문제

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

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

 

아이디어

y=ax+b 꼴이므로 입력값의 개수가 세 개 이상이면 a,b를 구하여 예측값과 비교를 하면 된다.

입력값의 개수가 한개라면 A를 출력

만약 두개라면 공차가 0인 등차수열인지 (두개의 값이 같다면) 값을 출력해주고, 다르다면 등차수열인지 등비수열인지 알수 없으므로 B를 출력한다.

 

코드

n=int(input())
arr=list(map(int, input().split()))

if n==1: print('A')
elif n==2:
    if arr[0]==arr[1]: print(arr[0])
    else: print('A')
else:
    # y = ax + b , 기울기랑 상수값 찾자
    if arr[1]-arr[0]!=0:
        a=(arr[2]-arr[1]) // (arr[1]-arr[0])
    else: 
        a=0

    b=arr[1]-a*arr[0]
    flag=0
    for j in range(n-1):
        tmp=a*arr[j]+b
        if arr[j+1]==tmp:
            continue
        else:
            print('B')
            flag=1
            break
    
    if flag!=1:
        print(arr[n-1]*a+b)