💻 문제: IQ Test

  • 알고리즘 분류: 구현(Implementation), 수학(Math)

문제 소개 🧐

수열이 주어졌을 때, 다음 수 = 앞 수 * a + b 라는 규칙을 찾아 다음 수를 예측하는 문제다. a와 b는 정수여야 한다. 규칙에 맞는 다음 수가 하나로 정해지면 그 수를, 여러 개일 수 있으면 ‘A’를, 규칙을 찾을 수 없으면 ‘B’를 출력해야 한다.

시간 제한 메모리 제한 제출 정답 맞힌 사람 정답 비율
2 초 128 MB 17981 3895 3165 22.245%

입력 📝

  • 첫째 줄에 수열의 크기 N (1 ≤ N ≤ 50)이 주어진다.
  • 둘째 줄에는 N개의 수가 주어진다. 각 수는 절댓값이 100보다 작거나 같은 정수다.

출력 📤

  • 규칙에 맞는 다음 수를 출력한다.
  • 다음 수가 여러 개일 경우 ‘A’를 출력한다.
  • 다음 수를 구할 수 없는 경우 ‘B’를 출력한다.

풀이: 경우의 수에 따른 구현 💪

Idea

이 문제는 주어진 수열에서 S[i] = S[i-1] * a + b를 만족하는 정수 a와 b를 찾는 것이 핵심이다. 여러 예외 케이스를 차분히 나누어 생각해야 했다.

  1. N = 1일 때: 숫자가 하나만 주어지면 규칙을 특정할 수 없다. 어떤 a, b를 사용해도 되므로 다음 수는 무한히 많다. 따라서 ‘A’를 출력한다.
  2. N = 2일 때:
    • 두 수가 같다면 (S[0] == S[1]), a=1, b=0인 등차수열일 수도 있고, 다른 여러 규칙이 가능하지만 가장 단순한 규칙은 모든 수가 같은 경우다. 따라서 다음 수는 S[1]과 같은 수가 된다.
    • 두 수가 다르다면, S[1] = S[0] * a + b를 만족하는 정수 (a, b) 쌍이 여러 개 존재하므로 규칙을 하나로 특정할 수 없다. 따라서 ‘A’를 출력한다.
  3. N > 2일 때:
    • S[1] = S[0] * a + b
    • S[2] = S[1] * a + b
    • 위 두 식을 연립하여 a를 구하면 a = (S[2] - S[1]) / (S[1] - S[0])가 된다.
    • 만약 S[1] - S[0]이 0이라면, 0으로 나눌 수 없으므로 별도 처리가 필요하다. 이때 모든 수가 같다면 다음 수도 그 수와 같지만, 다른 수가 하나라도 있다면 규칙이 성립하지 않는다.
    • a가 정수가 아니면 규칙이 성립하지 않으므로 ‘B’를 출력한다.
    • a를 구한 뒤 b = S[1] - S[0] * a로 b를 구한다.
    • 찾아낸 a와 b가 수열의 나머지 모든 수에 대해서도 성립하는지 검증한다. 만약 성립하지 않는 구간이 있다면 ‘B’를 출력한다.
    • 모든 검증을 통과하면 S[N-1] * a + b를 계산하여 다음 수를 출력한다.

이러한 케이스 분리를 바탕으로 코드를 작성했다.

Code

import sys

N = int(sys.stdin.readline().rstrip())
num_list = list(map(int, sys.stdin.readline().split()))

# 1. 각 숫자의 차를 구함
num_diff_list = list(map(lambda x, y: y - x, num_list[:-1], num_list[1:]))

def get_next_num():
    """
    다음 값을 구하는 함수
    다음 값이 여러개면 A, 없으면 B
    """
    global N
    global num_list
    global num_diff_list
    answer = 0

    # 2. len(num_diff_list) 가 1이나 0 밖에 없다면 규칙을 못 찾음 -> 암거나 넣어도 됨 A, 근데 0 하나 들어 있으면 첫번째 숫자 그대로
    if len(num_diff_list) < 2:
        if num_diff_list:
            return 'A' if num_diff_list[0] != 0 else num_list[0]
        else:
            return 'A'
    
    # 3. (앞수)*a + b 구함
    # 3-1. a 구하기, 0으로 나누면 안되니까 예외처리
    a = (num_diff_list[1] // num_diff_list[0]) if num_diff_list[0] != 0 else 0

    # 3-2. b 구하기, a가 0인 경우 b 만으로 반복이 이어져야 하므로 list 끝까지 검증해봄
    b = num_list[1] - num_list[0] * a
    for i in range(2, N):
        cur_b = num_list[i] - num_list[i-1] * a
        if b != cur_b:
            return 'B'

    # 4. 공식에 따라 output 구하고, 정수가 아니라면 B 출력
    answer = num_list[-1] * a + b
    if answer.is_integer():
        return int(answer)
    else:
        return 'B'

print(get_next_num())

Result

  • 결과: 성공! 🎉

배운 점 및 개선할 부분 🤔

  • 엣지 케이스의 중요성: N이 1 또는 2일 때와 같이 작은 입력값에 대한 처리가 문제의 핵심이었다. 이런 구현 문제는 항상 예외적인 상황을 먼저, 그리고 꼼꼼히 생각하는 습관이 중요하다.
  • 논리적 전개: ‘규칙을 찾을 수 없다’와 ‘규칙이 여러 개다’는 다른 상황이다. 문제가 요구하는 출력 조건(‘A’ 또는 ‘B’)을 명확히 구분하고 각 상황에 맞게 코드를 분기해야 한다.
  • 연립방정식: 간단한 연립방정식을 통해 규칙(a, b)을 유도하고, 그 규칙이 일반적인 경우에도 성립하는지 확인하는 과정은 다른 문제에서도 유용하게 쓰일 수 있는 접근법이다.

복잡해 보이는 문제일수록 단순한 케이스부터 차근차근 정복해나가자.

⁂