Computer >> 컴퓨터 >  >> 프로그래밍 >> C++

C++로 해결하는 등차 수열 슬라이스(Arithmetic Slices) 문제

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

이번 문제에서는 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]로 총 세 개의 등차 슬라이스가 존재하기 때문입니다.

해결 접근 방법

이 문제는 동적 계획법(Dynamic Programming)을 활용하면 효율적으로 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

  • ret := 0으로 초기화하고, n := 배열 A의 크기로 설정한 뒤, 크기 n의 dp 배열을 생성합니다.
  • i를 2부터 n-1까지 반복하며 다음을 검사합니다.
    • 만약 a[i] − a[i−1] = a[i−1] − a[i−2]라면, 즉 연속된 세 요소가 등차를 이룬다면:
      • dp[i] := 1 + dp[i−1]
      • ret에 dp[i]를 더합니다.
  • 최종적으로 ret을 반환합니다.

여기서 dp[i]는 '인덱스 i에서 끝나는 등차 슬라이스의 개수'를 의미합니다. 새로운 요소가 등차 조건을 만족할 때마다, 이전 인덱스에서 끝나던 모든 등차 슬라이스에 새 요소를 덧붙인 슬라이스들이 추가로 생성되므로 dp[i] = dp[i−1] + 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(){
    vector<int> v = {1,2,3,4};
    Solution ob;
    cout << (ob.numberOfArithmeticSlices(v));
}

입력

[1,2,3,4]

출력

3

복잡도 분석

시간 복잡도는 O(n)으로, 배열을 한 번만 순회하면 충분합니다. 공간 복잡도는 dp 배열 사용 시 O(n)이지만, 직전 값(dp[i−1])만 참조하므로 변수 하나만 유지하면 O(1)까지 최적화할 수 있습니다.