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

C++로 합이 K로 나누어 떨어지는 부분 배열 개수 구하기

정수로 이루어진 배열 A가 주어졌을 때, 원소들의 합이 k로 나누어 떨어지는 연속된 비어 있지 않은 부분 배열의 개수를 구하는 문제입니다.

예를 들어 A = [4, 5, 0, -2, -3, 1]이고 k = 5라고 가정해 보겠습니다. 이 경우 정답은 7이며, 조건을 만족하는 부분 배열은 다음과 같습니다.

  • [4, 5, 0, -2, -3, 1]
  • [5]
  • [5, 0]
  • [5, 0, -2, -3]
  • [0]
  • [0, -2, -3]
  • [-2, -3]

접근 방법: 누적 합과 나머지 연산

이 문제는 누적 합(prefix sum)과 나머지(modulo) 연산을 활용하면 O(n) 시간 복잡도로 효율적으로 해결할 수 있습니다.

핵심 아이디어는 다음과 같습니다. 두 누적 합을 k로 나눈 나머지가 서로 같다면, 그 사이에 있는 부분 배열의 합은 반드시 k로 나누어 떨어집니다. 따라서 각 나머지 값이 지금까지 몇 번 등장했는지 해시 맵에 기록하고, 현재 나머지와 같은 값을 가진 이전 누적 합의 개수만큼 정답에 더해주면 됩니다.

알고리즘 단계

  1. 맵 m을 생성하고 m[0]을 1로 초기화합니다. (나머지가 0인 경우, 즉 배열의 시작부터 해당 위치까지의 부분 배열을 처리하기 위함입니다.)
  2. temp := 0, ans := 0으로 초기화하고, n을 배열 a의 크기로 설정합니다.
  3. i를 0부터 n-1까지 순회하며 다음을 수행합니다:
    • temp := temp + a[i] → 누적 합을 계산합니다.
    • x := (temp mod k + k) mod k → 음수 처리를 포함한 나머지를 계산합니다.
    • ans := ans + m[x] → 같은 나머지를 가진 이전 누적 합의 개수만큼 정답에 더합니다.
    • m[x]를 1 증가시킵니다.
  4. 최종적으로 ans를 반환합니다.

(temp % k + k) % k처럼 나머지 연산을 두 번 수행하는 이유는, C++에서 음수에 대한 나머지 연산 결과가 음수가 될 수 있기 때문입니다. 예를 들어 temp가 -3이고 k가 5라면 -3 % 5는 -3이 되지만, 실제로 필요한 값은 2입니다.

C++ 구현 예제

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
    int subarraysDivByK(vector<int>& a, int k) {
        unordered_map <int, int> m;
        m[0] = 1;
        int temp = 0;
        int ans = 0;
        int n = a.size();
        for(int i = 0; i < n; i++){
            temp += a[i];
            int x = (temp % k + k) % k;
            ans += m[x];
            m[x]++;
        }
        return ans;
    }
};
main(){
    vector<int> v = {4,5,0,-2,-3,1};
    Solution ob;
    cout <<(ob.subarraysDivByK(v, 5));
}

입력

[4,5,0,-2,-3,1]
5

출력

7

복잡도 분석

배열을 한 번만 순회하므로 시간 복잡도는 O(n)입니다. 또한 해시 맵에는 최대 k개의 서로 다른 나머지만 저장되므로, 공간 복잡도는 O(min(n, k))입니다. 완전 탐색으로 모든 부분 배열을 확인하는 O(n²) 방식보다 훨씬 효율적이라는 점이 이 알고리즘의 가장 큰 장점입니다.