파이썬으로 숫자 목록에서 산술(등차) 부분 수열의 개수 찾는 프로그램

숫자 목록 nums가 주어졌을 때, 길이가 3 이상인 산술(arithmetic) 부분 수열의 개수를 구해야 합니다. 여기서 산술 수열(등차수열)이란 인접한 두 숫자 사이의 차이가 항상 일정한 수열을 의미합니다.

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

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

문제 해결 접근 방법

이 문제는 동적 계획법(Dynamic Programming)을 활용하면 효율적으로 해결할 수 있습니다. 각 인덱스와 공차(diff) 조합별로 만들 수 있는 부분 수열의 개수를 딕셔너리에 저장하면서 누적해 나가는 방식입니다.

구체적인 알고리즘 단계는 다음과 같습니다.

• dp : (인덱스, 공차)를 키로 갖는 새로운 맵(딕셔너리) 생성
• n : nums의 크기 저장
• res : 결과값을 담을 변수, 초기값은 0
• i를 0부터 n-1까지 반복:
  • j를 0부터 i-1까지 반복:
    • diff := nums[i] - nums[j] (두 원소의 공차 계산)
    • prev := dp[(i, diff)], 키가 없으면 0
    • prevprev := dp[(j, diff)], 키가 없으면 0
    • dp[(i, diff)] := prev + prevprev + 1
    • res := res + prevprev (길이 3 이상인 수열만 결과에 누적)
• 모든 반복이 끝나면 res 반환

여기서 prevprev는 j 위치에서 같은 공차로 형성된 기존 수열의 개수를 의미하며, 이 값이 곧 길이 3 이상으로 확장되는 수열의 수를 나타냅니다.

예제 코드

class Solution:
    def solve(self, nums):
        dp = {}
        n = len(nums)
        res = 0
        for i in range(n):
            for j in range(i):
                diff = nums[i] - nums[j]

                prev = dp.get((i, diff), 0)
                prevprev = dp.get((j, diff), 0)
                dp[(i, diff)] = prev + prevprev + 1

                res += prevprev
        return res

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

입력

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

출력

3

복잡도 분석

이 알고리즘은 모든 원소 쌍 (i, j)를 한 번씩 확인하므로 시간 복잡도는 O(n²)이며, 딕셔너리에 상태를 저장하기 때문에 공간 복잡도 역시 최대 O(n²)입니다. 원소의 개수가 수천 개 수준이라면 충분히 실용적인 성능을 보여줍니다.