C++ – 부분 배열 원소의 평균이 나머지 원소의 평균보다 큰 부분 배열의 개수 구하기

양의 정수로 이루어진 배열 arr[ ]가 주어졌을 때, 우리의 목표는 arr[ ]의 모든 부분 배열(subarray) 중에서 해당 부분 배열에 포함된 원소들의 평균이 포함되지 않은 나머지 원소들의 평균보다 큰 부분 배열의 개수를 구하는 것입니다.

예제

입력

arr[ ] = { 3, 2, 4 }

출력

부분 배열에 포함된 원소의 평균이 포함되지 않은 원소의 평균보다 큰 부분 배열의 개수: 2

설명

가능한 부분 배열은 다음과 같습니다.
[ 3 ], [ 2 ], [ 4 ], [ 3,2 ], [ 2,4 ], [ 3,2,4 ]
[ 4 ]의 평균은 4로, 나머지 원소들 [ 2,3 ]의 평균 2.5보다 큽니다.
[ 3,2,4 ]의 평균은 3으로, 남는 원소가 없는 빈 집합의 평균(0)보다 큽니다.

입력

arr[ ] = { 3, 3, 3 }

출력

부분 배열에 포함된 원소의 평균이 포함되지 않은 원소의 평균보다 큰 부분 배열의 개수: 1

설명

가능한 부분 배열은 다음과 같습니다.
[ 3 ], [ 3 ], [ 3 ], [ 3,3 ], [ 3,3 ], [ 3,3,3 ]
모든 원소가 3으로 동일하므로, 전체 배열 [ 3,3,3 ](평균 3)이 빈 집합의 평균(0)보다 큰 유일한 경우입니다.

알고리즘(접근 방식)

이 문제는 접두사 합(prefix sum) 배열을 이용해 효율적으로 해결할 수 있습니다. 인덱스 i까지의 원소 합을 new_arr[i]에 미리 저장해 두면, 임의의 구간 [i, j]의 합을 상수 시간에 구할 수 있고, 구간에 포함된 원소의 개수는 j-i+1이 되므로 각 구간의 평균을 손쉽게 계산할 수 있습니다.

배열 arr[ ]를 입력으로 받습니다.
count(int arr[], int size) 함수는 arr[ ]를 받아, 부분 배열에 포함된 원소의 평균이 포함되지 않은 원소의 평균보다 큰 부분 배열의 개수를 반환합니다.
크기가 size+1인 배열 new_arr[ ]를 선언하고, i=1부터 i<size까지 new_arr[i]를 new_arr[i-1] + arr[i-1]로 채워 접두사 합을 구성합니다.
두 개의 for 루프를 사용해 가능한 모든 구간 [i, j]를 탐색합니다.
total_1은 현재 구간의 합, count_1은 해당 구간에 포함된 원소의 개수입니다.
total_2는 현재 구간 외 나머지 부분의 합, count_2는 그 원소의 개수입니다. 구간이 배열 전체인 경우 0으로 나누는 오류를 피하기 위해 count_2를 1로 설정합니다.
평균을 check_1 = total_1 / count_1, check_2 = total_2 / count_2로 계산합니다.
check_1 > check_2를 만족하면 count를 증가시킵니다.
모든 루프가 끝나면 count를 결과로 반환합니다.

C++ 구현 예제

#include <bits/stdc++.h>
using namespace std;

int count(int arr[], int size){
   int count = 0;
   int new_size = size + 1;
   int new_arr[new_size] = { 0 };
   for (int i = 1; i < new_size; i++){
      new_arr[i] = new_arr[i - 1] + arr[i - 1];
   }
   for (int i = 1; i < new_size; i++){
      for (int j = i; j < new_size; j++){
         int total_1 = new_arr[j] - new_arr[i - 1];
         int count_1 = j - i + 1;
         int total_2 = new_arr[size] - total_1;
         int count_2 = 0;
         if((size - count_1) == 0){
            count_2 = 1;
         } else {
            count_2 = size - count_1;
         }
         int check_1 = total_1 / count_1;
         int check_2 = total_2 / count_2;
         if (check_1 > check_2){
            count++;
         }
      }
   }
   return count;
}

int main(){
   int arr[] = { 2, 6, 2, 4 };
   int size = sizeof(arr) / sizeof(arr[0]);
   cout << "부분 배열에 포함된 원소의 평균이 포함되지 않은 원소의 "
        << "평균보다 큰 부분 배열의 개수: " << count(arr, size);
   return 0;
}

출력

위 코드를 실행하면 다음과 같은 결과가 출력됩니다.

부분 배열에 포함된 원소의 평균이 포함되지 않은 원소의 평균보다 큰 부분 배열의 개수: 6

복잡도 분석

모든 구간 [i, j]를 두 개의 중첩 루프로 탐색하므로 시간 복잡도는 O(N²)이며, 접두사 합 배열을 저장하기 위해 O(N)의 추가 공간이 필요합니다. 또한 이 구현은 정수 나눗셈으로 평균을 비교하기 때문에, 실수 기반 평균 비교를 사용했을 때와 결과가 달라질 수 있다는 점을 유의해야 합니다.