파이썬으로 숫자 리스트에서 등차수열 개수 찾는 방법

숫자 리스트 nums가 주어졌을 때, 길이가 3 이상인 연속된 등차수열(arithmetic sequence)의 개수를 구하는 문제입니다. 등차수열이란 인접한 두 숫자 사이의 차이가 항상 일정한 수열을 의미합니다.

예를 들어 입력이 nums = [6, 8, 10, 12, 13, 14]라면 출력은 4가 됩니다. 다음과 같은 등차수열들이 존재하기 때문입니다.

• [6, 8, 10]
• [8, 10, 12]
• [6, 8, 10, 12]
• [12, 13, 14]

문제 해결 접근 방식

이 문제는 한 번의 순회(O(n))로 효율적으로 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

1. count와 ans를 0으로 초기화합니다.
2. 인덱스 2부터 리스트 끝까지 반복하면서, 현재 요소와 이전 요소의 차이가 그 앞 두 요소의 차이와 같은지 확인합니다.
3. 차이가 같으면 count를 1 증가시켜 등차수열이 계속 이어지고 있음을 기록합니다.
4. 차이가 달라지면 지금까지 누적된 count로 만들 수 있는 부분 수열의 개수를 정답에 더하고, count를 0으로 초기화합니다.
5. 반복이 끝난 후 count가 남아 있다면 마지막으로 한 번 더 더해줍니다.

왜 count × (count + 1) ÷ 2일까?

연속된 차이가 count번 일정하다면, 길이 3 이상인 등차수열은 count * (count + 1) // 2개가 됩니다. 예를 들어 차이가 3번 연속 같으면(count = 3), 길이 3짜리 수열 3개, 길이 4짜리 2개, 길이 5짜리 1개로 총 6개 = 3 × 4 ÷ 2가 됩니다. 이는 1부터 count까지의 합 공식과 동일합니다.

구현 예제

class Solution:
    def solve(self, nums):
        count = 0
        ans = 0
        for i in range(2, len(nums)):
            if nums[i] - nums[i - 1] == nums[i - 1] - nums[i - 2]:
                count += 1
            else:
                ans += (count * (count + 1)) // 2
                count = 0
        if count:
            ans += (count * (count + 1)) // 2
        return ans

ob = Solution()
nums = [6, 8, 10, 12, 13, 14]
print(ob.solve(nums))

입력

[6, 8, 10, 12, 13, 14]

출력

4

복잡도 분석

이 알고리즘은 리스트를 한 번만 순회하므로 시간 복잡도는 O(n)이며, 추가 변수 몇 개만 사용하므로 공간 복잡도는 O(1)입니다. 따라서 매우 큰 입력 데이터에서도 효율적으로 동작합니다.