숫자 리스트 nums와 쿼리 리스트 queries가 주어져 있다고 가정해 봅시다. 각 쿼리는 [i, j] 형태로 구성되며, 해당 쿼리는 nums의 i번째부터 j번째까지(양 끝 인덱스 포함) 부분 리스트가 산술 수열(arithmetic sequence)인지를 묻습니다. 우리의 목표는 참(true)을 반환하는 쿼리의 개수를 구하는 것입니다.
예제로 이해하기
예를 들어 입력이 다음과 같다고 해보겠습니다.
- nums = [2, 4, 6, 8, 7, 6, 5, 2]
- queries = [[3, 4], [0, 3], [2, 4]]
이때 출력은 2가 됩니다.
- [2, 4, 6, 8]은 공차가 2로 일정한 산술 수열이므로 쿼리 [0, 3]은 참입니다.
- [8, 7]은 두 원소만으로 이루어진 시퀀스로 항상 산술 수열이 성립하므로 쿼리 [3, 4]도 참입니다.
- 반면 [6, 8, 7]은 연속된 원소 간 차이가 일정하지 않으므로 쿼리 [2, 4]는 거짓입니다.
해결 접근 방법
쿼리마다 매번 부분 배열을 처음부터 검사하면 비효율적입니다. 대신 인접한 두 원소의 차이(diff)를 미리 계산하고, 동일한 차이가 연속으로 이어지는 길이(rle, run-length)를 저장해 두면 각 쿼리를 상수 시간 O(1)에 판별할 수 있습니다.
알고리즘의 단계는 다음과 같습니다.
- nums가 비어 있다면 0을 반환합니다.
- n := nums의 크기로 설정합니다.
- diff := 각 i(0 ~ n-2)에 대해 nums[i + 1] - nums[i] 값을 담은 리스트를 만듭니다.
- rle := 크기가 n - 1이고 0으로 초기화된 리스트를 만듭니다.
- i를 0부터 n - 2까지 반복합니다.
- i > 0이고 diff[i] == diff[i - 1]이면 rle[i] := rle[i - 1] + 1
- 그렇지 않으면 rle[i] := 1 - ans := 0으로 초기화합니다.
- 각 쿼리 (i, j)에 대해 다음을 수행합니다.
- i == j이면 ans를 1 증가시킵니다. (원소가 하나뿐인 시퀀스는 항상 산술 수열입니다.)
- 그렇지 않으면 rle[j - 1] >= (j - i)를 만족할 때 ans를 1 증가시킵니다. - ans를 반환합니다.
Python 구현 코드
아래 구현을 통해 더 자세히 이해해 보겠습니다.
def solve(nums, queries):
if not nums:
return 0
n = len(nums)
diff = [nums[i + 1] - nums[i] for i in range(n - 1)]
rle = [0] * (n - 1)
for i in range(n - 1):
if i > 0 and diff[i] == diff[i - 1]:
rle[i] = rle[i - 1] + 1
else:
rle[i] = 1
ans = 0
for i, j in queries:
if i == j:
ans += 1
else:
ans += rle[j - 1] >= (j - i)
return ans
nums = [2, 4, 6, 8, 7, 6, 5, 2]
queries = [[3, 4],[0, 3],[2, 4]]
print(solve(nums, queries))
입력
[2, 4, 6, 8, 7, 6, 5, 2], [[3, 4],[0, 3],[2, 4]]
출력
2
동작 원리 정리
rle 배열은 "동일한 공차가 몇 번 연속으로 나타나는지"를 누적으로 기록합니다. 구간 [i, j]가 산술 수열이 되려면 i부터 j-1까지의 모든 차이가 같아야 하며, 이는 rle[j - 1]이 (j - i) 이상이라는 조건으로 손쉽게 확인할 수 있습니다. 덕분에 전처리 이후에는 각 쿼리를 상수 시간에 처리할 수 있어, 쿼리 개수가 많은 경우에도 매우 효율적으로 동작합니다.