알고리즘 이론

나머지 부분합 배열 및 나머지가 같은 배열 개수 구하기

행수쌤 2023. 3. 8. 15:49
728x90
반응형

나머지 부분합 배열

나머지 부분합 배열(Remainder Prefix Sum Array)은 주어진 수열의 각 원소를 어떤 수 K로 나눈 나머지를 누적한 배열을 의미합니다.

이를 계산하기 위해서는 부분합 배열(Prefix Sum Array)을 먼저 계산해야 합니다. 부분합 배열이란, 주어진 수열의 각 원소를 누적한 배열로, 예를 들어 주어진 수열이 [1, 2, 3, 4, 5]라면 부분합 배열은 [1, 3, 6, 10, 15]가 됩니다.

나머지 부분합 배열은 이 부분합 배열의 각 원소를 K로 나눈 나머지를 누적한 배열입니다. 예를 들어, 주어진 수열이 [1, 2, 3, 4, 5]이고 K가 3이라면, 부분합 배열은 [1, 3, 6, 10, 15]가 되고, 이를 3으로 나눈 나머지를 누적한 나머지 부분합 배열은 [1, 0, 0, 1, 2]가 됩니다.

이를 자바 코드로 구현하면 다음과 같습니다.

public int[] remainderPrefixSum(int[] arr, int K) {
    int[] prefixSum = new int[arr.length];
    int[] remainderPrefixSum = new int[arr.length];

    prefixSum[0] = arr[0];
    remainderPrefixSum[0] = arr[0] % K;

    for (int i = 1; i < arr.length; i++) {
        prefixSum[i] = prefixSum[i - 1] + arr[i];
        remainderPrefixSum[i] = prefixSum[i] % K;
    }

    return remainderPrefixSum;
}

이 코드는 입력으로 주어진 배열 arr과 나누는 수 K를 받아 나머지 부분합 배열을 계산하여 반환합니다.

먼저, 배열 prefixSum과 remainderPrefixSum을 초기화하고, prefixSum[0]과 remainderPrefixSum[0]은 arr[0]의 값으로 초기화합니다.

그리고 반복문을 이용하여 prefixSum 배열을 계산하고, 이를 이용하여 remainderPrefixSum 배열을 계산합니다.

prefixSum[i]는 prefixSum[i-1]과 arr[i]를 더한 값이므로, remainderPrefixSum[i]는 prefixSum[i]를 K로 나눈 나머지가 됩니다.

마지막으로, 계산된 remainderPrefixSum 배열을 반환합니다.

 

부분합의 나머지가 같은 배열 개수 구하기

 

주어진 수열의 부분합 중에서 k로 나누어 떨어지는 부분합의 개수를 구하는 방법은 나머지 부분합 배열(Remainder Prefix Sum Array)을 이용하는 것입니다.

나머지 부분합 배열에서 같은 나머지 값을 가지는 원소들을 두 개씩 선택하면, 그 차이는 k로 나누어 떨어지는 수가 됩니다. 따라서 같은 나머지 값을 가지는 원소들의 개수를 구하면, 이들을 두 개씩 선택하여 나머지가 0이 되는 부분합의 개수를 구할 수 있습니다.

이를 자바 코드로 구현하면 다음과 같습니다.

public int countSubarraysDivisibleByK(int[] arr, int K) {
    int[] remainderPrefixSum = remainderPrefixSum(arr, K);
    int[] count = new int[K];

    for (int i = 0; i < remainderPrefixSum.length; i++) {
        count[remainderPrefixSum[i]]++;
    }

    int result = count[0];
    for (int i = 0; i < K; i++) {
        result += count[i] * (count[i] - 1) / 2;
    }

    return result;
}

이 코드는 입력으로 주어진 배열 arr과 나누는 수 K를 받아, k로 나누어 떨어지는 부분합의 개수를 반환합니다.

먼저, 앞서 구현한 나머지 부분합 배열 함수를 이용하여 나머지 부분합 배열 remainderPrefixSum을 계산합니다.

그리고 나머지 값이 0부터 K-1까지 나오는 원소들의 개수를 저장하는 count 배열을 초기화합니다.

이후, remainderPrefixSum 배열을 탐색하면서 각 나머지 값을 가지는 원소들의 개수를 count 배열에 저장합니다.

마지막으로, count[0]은 자기 자신만으로도 k로 나누어 떨어지는 부분합의 개수를 의미하므로, 이를 먼저 result에 더해줍니다. 그리고 1부터 K-1까지의 나머지 값을 가지는 원소들 중에서 두 개씩 선택하여 나머지가 0이 되는 부분합의 개수를 계산하여 result에 더해줍니다.

이렇게 계산된 result를 반환합니다.

 

for (int i = 0; i < K; i++) {
        result += count[i] * (count[i] - 1) / 2;
    }

이 부분은 나머지 부분합 배열에서 같은 나머지 값을 가지는 원소들을 두 개씩 선택하여 그 차이가 k로 나누어 떨어지는 부분합의 개수를 구하는 것을 의미합니다.

count 배열에는 0부터 K-1까지의 나머지 값들이 각각 몇 개씩 나오는지 저장되어 있습니다. 따라서, count[i]는 나머지가 i인 원소의 개수를 의미합니다.

이제, 나머지가 i인 원소들 중에서 두 개씩 선택하는 경우의 수는 count[i]개 중에서 2개를 선택하는 경우의 수와 같습니다. 이 경우의 수는 nC2로 구할 수 있습니다. 즉, count[i]개 중에서 2개를 선택하는 경우의 수는 count[i] * (count[i] - 1) / 2입니다.

그리고 이렇게 구한 경우의 수를 모두 더하면, 나머지 부분합 배열에서 같은 나머지 값을 가지는 원소들을 두 개씩 선택하여 그 차이가 k로 나누어 떨어지는 부분합의 총 개수가 됩니다. 따라서, result 변수에 이 값을 누적하여 저장합니다.

728x90
반응형