[프로그래머스] 다음에 올 숫자 / Python

문제

https://school.programmers.co.kr/learn/courses/30/lessons/120924

프로그래머스

등차수열 혹은 등비수열 common이 매개변수로 주어질 때, 마지막 원소 다음으로 올 숫자를 return 하도록 solution 함수를 완성해보세요.

 

제한사항

• 2 < common의 길이 < 1,000
• -1,000 < common의 원소 < 2,000
• common의 원소는 모두 정수입니다.
• 등차수열 혹은 등비수열이 아닌 경우는 없습니다.
• 등비수열인 경우 공비는 0이 아닌 정수입니다.

 

아이디어

등차수열이라면 1,2항의 차와 마지막 끝에 2항의 차가 같을 것이고, 아니라면 등비 수열일 것이다.

마지막 원소에 공차 만큼 더하거나, 공비 만큼 곱하면 답이 나온다.

 

풀이

import math
def solution(common):
    if common[1] - common[0] == common[-1]-common[-2]:
        return common[-1] + (common[1] - common[0])
    else:
        return common[-1]*(common[1]/common[0])