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

산술 슬라이스 II – C++로 구하는 등차 부분 수열의 개수

문제 설명

N개의 숫자로 이루어진 배열 A가 있다고 가정해 봅시다. 배열 A의 하위 시퀀스 슬라이스(subsequence slice)란 (K0, K1, K2, …, Kn)과 같이 0 ≤ K0 < K1 < K2 < … < Kn < N을 만족하는 정수 인덱스의 나열을 의미합니다. 이때 수열 A[K0], A[K1], …, A[Kn]이 등차수열을 이룬다면, 즉 항의 개수가 최소 3개(n ≥ 2)라면 이 하위 시퀀스 슬라이스를 산술 슬라이스(arithmetic slice)라고 부릅니다. 따라서 우리가 구해야 하는 값은 배열 A에서 만들 수 있는 산술 슬라이스의 총 개수입니다.

예를 들어 입력이 [2, 4, 6, 8, 10]이라면 정답은 7이며, 만들어지는 산술 슬라이스는 다음과 같습니다.

  • [2, 4, 6]
  • [4, 6, 8]
  • [6, 8, 10]
  • [2, 6, 10]
  • [2, 4, 6, 8]
  • [4, 6, 8, 10]
  • [2, 4, 6, 8, 10]

접근 방법 – 동적 계획법(DP)

이 문제는 동적 계획법을 활용해 효율적으로 해결할 수 있습니다. 핵심 아이디어는 dp[i][diff]를 “인덱스 i에서 끝나고 공차가 diff인 등차 부분 수열의 개수”로 정의하는 것입니다. 그러면 각 쌍 (j, i)를 검사하면서 이전 결과를 확장해 나갈 수 있습니다.

구체적인 알고리즘 진행 과정은 다음과 같습니다.

  • ret := 0 으로 초기화합니다.
  • dp 맵과 cnt 맵을 하나씩 정의합니다.
  • A의 원소들을 모두 포함하는 집합 s를 만듭니다.
  • n := A의 크기로 설정합니다.
  • i := 1부터 i < n일 때까지 1씩 증가시키며 반복합니다.
    • j := i − 1부터 j ≥ 0일 때까지 1씩 감소시키며 반복합니다.
      • diff := A[i] − A[j] 로 공차를 계산합니다.
      • diff가 int 범위(INT_MIN ~ INT_MAX)를 벗어나면 다음 반복으로 건너뜁니다.
      • temp := dp[j]에 diff가 존재하면 dp[j][diff], 그렇지 않으면 0
      • ret := ret + temp
      • A[i] + diff가 집합 s에 존재한다면, dp[i][diff] := dp[i][diff] + temp + 1
  • 최종적으로 ret을 반환합니다.

여기서 temp는 인덱스 j에서 끝나던 기존 등차 수열의 개수입니다. 이 수열들 뒤에 A[i]를 붙이면 길이가 3 이상인 완성된 산술 슬라이스가 되므로 정답(ret)에 더해집니다. 반면 +1에 해당하는 새로운 쌍 (A[j], A[i])은 아직 길이가 2에 불과하므로 ret에는 포함하지 않고, 이후 확장을 위해 dp[i][diff]에만 누적합니다. 집합 s를 활용해 A[i] + diff 값이 실제로 배열에 존재할 때만 갱신하면 불필요한 연산을 줄일 수 있습니다. 전체 시간 복잡도는 O(N²)입니다.

아래 예제 코드를 통해 더 자세히 이해해 보겠습니다.

C++ 구현 예제

#include <bits/stdc++.h>
using namespace std;
typedef long long int lli;
class Solution {
public:
    int numberOfArithmeticSlices(vector<int>& A) {
        int ret = 0;
        unordered_map<lli, unordered_map<lli, lli>> dp, cnt;
        unordered_set<int> s(A.begin(), A.end());
        int n = A.size();
        for(int i = 1; i < n; i++){
            for(int j = i - 1; j >= 0; j--){
                lli diff = (lli)A[i] - (lli)A[j];
                if(diff <= INT_MIN || diff > INT_MAX) continue;
                int temp = dp[j].count(diff) ? dp[j][diff] : 0;
                ret += temp;
                if(s.count(A[i] + diff)) dp[i][diff] += temp + 1;
            }
        }
        return ret;
    }
};
main(){
    Solution ob;
    vector<int> v = {2,4,6,8,10};
    cout << (ob.numberOfArithmeticSlices(v));
}

입력

{2,4,6,8,10}

출력

7