C++ 프로그램으로 배열의 등차 슬라이스 개수 구하기

등차 수열이란 무엇인가?

등차 수열(arithmetic sequence)은 최소 세 개 이상의 요소로 구성되며, 인접한 두 요소 사이의 차이가 항상 일정한 수열을 의미합니다. 예를 들어 [1, 3, 5, 7, 9], [7, 7, 7, 7], [3, -1, -5, -9]는 모두 등차 수열입니다. 반면 [1, 1, 2, 5, 7]은 요소 간 차이가 0, 1, 3, 2로 일정하지 않기 때문에 등차 수열이 아닙니다.

문제 정의

N개의 숫자로 이루어진 0-인덱스 배열 A가 주어졌다고 가정해 보겠습니다. 배열의 슬라이스(slice)는 0 <= P < Q < N을 만족하는 정수 쌍 (P, Q)로 정의됩니다. 이때 슬라이스 (P, Q)에 해당하는 부분 수열 A[P], A[P + 1], ..., A[Q - 1], A[Q]가 등차 수열을 이룬다면, 이 슬라이스를 '등차 슬라이스(arithmetic slice)'라고 부릅니다.
목표는 배열 A 안에 존재하는 등차 슬라이스의 총개수를 구하는 함수를 작성하는 것입니다.

예를 들어 입력 배열이 [1, 2, 3, 4]라면 결과는 3입니다. 등차 슬라이스가 정확히 [1, 2, 3], [2, 3, 4], [1, 2, 3, 4] 세 개 존재하기 때문입니다.

풀이 접근 방법

이 문제는 동적 계획법(DP)으로 효율적으로 해결할 수 있습니다. dp[i]를 '인덱스 i에서 끝나는 등차 슬라이스의 개수'라고 정의하면, 연속한 세 요소 A[i-2], A[i-1], A[i]가 등차 관계를 만족할 때마다 새로운 슬라이스들이 생겨납니다. 구체적인 단계는 다음과 같습니다.

• 결과 변수 ret := 0, 배열 크기 n := size of A로 초기화하고, 크기가 n인 dp 배열을 생성합니다.
• i를 2부터 n - 1까지 반복합니다.
  • 만약 A[i] - A[i - 1] == A[i - 1] - A[i - 2]라면, 즉 연속된 세 요소가 등차 관계를 이룬다면:
    • dp[i] := 1 + dp[i - 1]로 갱신합니다. 이전 위치에서 끝나던 모든 등차 슬라이스를 A[i]까지 확장할 수 있고, 여기에 길이 3짜리 새로운 슬라이스 하나가 추가되기 때문입니다.
    • ret에 dp[i] 값을 누적합니다.
• 반복이 끝나면 ret을 반환합니다.

이 알고리즘의 시간 복잡도는 O(n), 공간 복잡도는 O(n)입니다. 참고로 dp 배열 전체 대신 바로 이전 값 하나만 저장하면 공간 복잡도를 O(1)로 줄일 수도 있습니다.

C++ 구현 예제

아래 구현을 통해 풀이 과정을 더 명확하게 이해할 수 있습니다.

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
    int numberOfArithmeticSlices(vector<int>& A) {
        int ret = 0;
        int n = A.size();
        vector<int> dp(n);
        for(int i = 2; i < n; i++){
            if(A[i] - A[i - 1] == A[i - 1] - A[i - 2]){
                dp[i] = 1 + dp[i - 1];
                ret += dp[i];
            }
        }
        return ret;
    }
};
main(){
    Solution ob;
    vector<int> v = {1, 2, 3, 4};
    cout << (ob.numberOfArithmeticSlices(v));
}

입력

[1, 2, 3, 4]

출력

3