[백준] 1111번 IQ Test (자바 풀이)

문제

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

1111번: IQ Test

풀이

처음 이 문제를 봤을 땐, a와 b를 -100부터 100까지만 판단해서 브루트 포스 방식으로 해결하려고 했다. 

하지만 a와 b를 -10000~10000까지 전부 조사해도 틀리고, 그 이상은 시간이 엄청나게 오래걸린다.

 

그런데 점화식을 잘 보면, a를 쉽게 구할 수 있음을 알 수 있다. 

1. arr[i+1] = arr[i]*a + b

2. arr[i] = arr[i-1]*a + b

 

(arr[i+1]-arr[i])  = (arr[i]*a+b) - (arr[i-1]*a+b) = (arr[i]-arr[i-1])*a  이다.

따라서, 

(arr[i+1]-arr[i]) / (arr[i]-arr[i-1]) = (arr[i]-arr[i-1])*a / (arr[i]-arr[i-1]) = a 이다.

 

a를 구하면, b는 쉽게 구할 수 있다. b= arr[i+1]-(arr[i]*a) 이기 때문이다.

 

이제 몇 개의 예외를 처리해주면 된다.

 

먼저, N=1 이면 무조건 답은 "A" 일수 밖에 없다. 다음에 올 수가 무수히 많기 때문이다.

또한 N=2 일때 arr[0]과 arr[1]이 다르다면 또한 무수한 규칙이 가능하므로 정답은 "A"다.

하지만 N=2 일때 arr[0]와 arr[1]이 같다면 다음 올 수는 무조건 같은 수이다. 

 

이러한 예외를 처리하고 나면, a와 b를 구해서 모든 원소가 규칙에 맞는지 확인하면 된다.

하지만 이때 주의할 점이 arr[1]과 arr[0]이 같으면 divide by 0 예외가 발생할 수 있으므로 따로 처리해주어야 한다.

코드

import java.util.*;
import java.io.*;
 
public class Baekjoon_1111 {
    static int N;
    static int arr[];
    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        N = sc.nextInt();
        arr = new int[N];
        for(int i=0; i<N; i++) {
            arr[i] = sc.nextInt();
        }
        
        if(N == 1 || (N==2 && arr[0] != arr[1])) { // 무수히 많은 규칙
            System.out.println("A");
            return;
        }else if(N==2) {
            System.out.println(arr[0]);
            return;
        }
        
        int a,b;
        if(arr[1] == arr[0]) {    // divide by 0 방지
            a = 1;
            b = 0;
        }else {
            a = (arr[2]-arr[1])/(arr[1]-arr[0]);
            b = arr[1]-(arr[0]*a);
        }
        
        if(check(a,b)) {
            System.out.println(arr[N-1]*a+b);
        }else {
            System.out.println("B");
        }
        
    }
    
    /* 모든 원소가 찾은 규칙을 만족하는지 체크 */
    static boolean check(int a, int b) {
        for(int i=0; i<N-1; i++) {
            if(arr[i+1] != (arr[i]*a+b) )
                return false;
        }
        return true;
    }
 
}

결과