파이썬으로 일정한 차이를 갖는 가장 긴 등차 부분 수열의 길이 구하기

숫자 리스트 nums와 정수 값 diff가 주어졌다고 가정해 봅시다. 이때 부분 수열 내에서 인접한 두 숫자의 차이가 항상 diff와 같도록 만들 수 있는 가장 긴 등차 부분 수열의 길이를 구하는 것이 목표입니다.

예를 들어 입력이 nums = [-1, 1, 4, 7, 2, 10], diff = 3이라면 출력은 4가 됩니다. 그 이유는 [1, 4, 7, 10]과 같은 부분 수열을 선택할 수 있고, 각 인접 원소의 차이가 모두 3으로 일정하기 때문입니다.

문제 해결 접근 방식

이 문제는 동적 계획법(DP)과 해시맵(딕셔너리)을 활용하면 O(n) 시간 복잡도로 효율적으로 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

• seen: 각 숫자를 마지막 원소로 하는 등차 부분 수열의 최대 길이를 저장하는 딕셔너리 (기본값은 0)
• mx: 지금까지 발견한 최대 부분 수열 길이
• 리스트의 각 원소 x에 대해 다음을 반복합니다.
  • x - diff가 seen에 존재하면, x를 마지막으로 하는 부분 수열의 길이는 seen[x - diff] + 1
  • 존재하지 않으면, x 자체로 시작하는 새로운 수열이므로 seen[x] = 1
  • mx를 seen[x]와 비교하여 갱신
• 모든 원소를 처리한 후 mx를 반환합니다.

구현 예제

아래 파이썬 코드를 통해 더 잘 이해해 보겠습니다.

from collections import defaultdict

def solve(nums, diff):
    seen = defaultdict(int)
    mx = 0
    for x in nums:
        if x - diff in seen:
            seen[x] = seen[x - diff] + 1
        else:
            seen[x] = 1
        mx = max(mx, seen[x])
    return mx

nums = [-1, 1, 4, 7, 2, 10]
diff = 3
print(solve(nums, diff))

입력

[-1, 1, 4, 7, 2, 10], 3

출력

4

복잡도 분석

리스트의 각 원소를 한 번씩만 순회하고 딕셔너리 조회는 평균 O(1)이므로, 시간 복잡도는 O(n), 공간 복잡도 역시 딕셔너리에 최대 n개의 키가 저장될 수 있으므로 O(n)입니다.