코드 포스 (Code Forces) - Good Prefix (C++)

 

 

{A1, A2, A3, .... An} 처럼 n개의 원소를 가진 배열이 있다고 해보자.

특정 원소를 Ak를 제외한 나머지 원소를 모두 더한 값이 Ak가 되는 Ak가 존재한다면, 해당 Array는 Good이라고 한다.

예를 들어, {1, 6, 3, 2}에서 6을 제외한 나머지 원소의 합은 6이 된다. 그러므로, 이 배열은 Good이라고 할 수 있다.

 

{A1, A2, A3, ..., An}과 같이 n개의 원소로 이루어진 배열이 있다고 해보자.

{A1}, {A1, A2}, {A1, A2, A3} ... 등 배열의 가장 앞에서부터 잘라낸 배열을 Prefix라고 할 것이다.

An에 대한 모든 Prefix중 Good인 Array가 몇 개인지를 구하여 출력하자.

 

문제 풀이

배열에서 특정 원소 Ak를 제외한 나머지 원소의 총 합이 Ak가 된다면, 해당 배열은 Good이라고 할 수 있다.

그렇다면, 해당 배열에서 Ak가 될 수 있는 후보는 어떤 숫자가 있을까?

 

Ak가 될 수 있는 값은, 해당 배열의 원소중 최대값이다.

 

{1, 2, 3, 6}이라는 배열이 있다고 해보자. 여기서, 최대값은 6이다. 이 배열에서 만약 Ak를 최대값이 아닌 다른 값으로 잡는다면, (Ak = 최대값 + 다른 원소) 의 공식이 성립해야 한다. 즉, Ak는 배열의 원소중 최대값보다 더 큰 값이어야 한다는 것인데 이는 모순되는 말이므로 Ak는 반드시 배열의 원소중 최대값이어야 한다.

 

즉, (배열의 최대값 = 나머지 원소의 총합)이 성립한다면 해당 배열은 Good이라고 할 수 있으며, 성립하지 않는다면 Good이 아니라고 할 수 있다.

 

그렇다면, 주어진 배열의 모든 Prefix 에 대해 Good을 판단할 수 있게 된다.

std::vector<int> Array;

for(int i = 0; i < Array.size(); i++)
{    
    int Max = 0;
    int Sum = 0;
    
    for(int j = 0; j <= i; j++)
    {
        Max = std::max(Max, Array[j]);
        Sum += Array[j];
    }
    
    if(Max - Sum == Sum)
    {
        //Array is Good!
    }
}

위와 같이 이중 반복문을 사용해 모든 Prefix 배열에 대해 Good을 판별할 수 있다.

하지만, 이렇게 무식하게 판별을 하게 되면 시간초과가 발생할 것이다.

그러므로, 판별 과정을 조금 더 단순하게 만들 필요가 있다.

 

결국 배열이 Good인지 판단하기 위해 중요한 것은 배열의 원소의 총 합과 최대값이다.

(배열의 원소의 총 합 - 최대값 = 최대값) 을 만족한다면 배열은 Good이라고 할 수 있기 때문이다.

 

생각해보면, Prefix 배열의 원소 총 합과 최대값을 구하기 위해서 매번 Prefix 배열의 사이즈만큼 반복문을 돌 필요는 없다.

모든 Prefix배열은 입력으로 주어진 가장 앞에서부터 시작하며, size가 K인 Prefix배열은 size가 (k - 1)인 배열의 가장 뒤에 Ak를 추가한 형태일 뿐이기 때문이다.

 

즉, size가 k인 배열의 원소 총 합은 size가 (k - 1)인 배열의 원소 총 합에 Ak를 더한 것이다.

또한 size가 k인 배열의 원소중 최대값은 size가 (k - 1)인 배열의 원소중 최대값과 Ak중 더 큰 값이 된다.

이를 코드로 구현하면 아래와 같이 간단하게 한 번의 반복문으로 표현할 수 있게 된다.

std::vector<int> Array;

int Sum = 0;
int Max = 0;

for(int i = 0; i < Array.size(); i++)
{    
    Max = std::max(Max, Array[i]);
    Sum += Array[i];
    
    if(Max - Sum == Sum)
    {
        //Prefix Array is Good!
    }
}

 

풀이 코드

int ArraySize = 0;
std::cin >> ArraySize;

std::vector<long long> Array(ArraySize);
for (int i = 0; i < ArraySize; i++)
{
    std::cin >> Array[i];
}

먼저 배열을 입력받아 준다.

int GoodCount = 0;

long long MaxNum = 0;
long long Sum = 0;

다음은 값을 저장할 변수들을 선언해주었다.

for (int i = 0; i < ArraySize; i++)
{
    Sum += Array[i];
    MaxNum = std::max(Array[i], MaxNum);

    if (MaxNum - Sum == Sum)
    {
        GoodCount++;
    }
}

앞에서 설명한 것과 같이 반복문을 돌며 모든 prefix에 대해 Good 여부를 판별한 뒤, Good이라면 GoodCount를 증가시켜주었다.

 

모든 케이스에 대해 GoodCount를 구한 뒤 출력해주면 된다.

 

코드 전문

#include <iostream>
#include <vector>

void Init()
{
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    std::cout.tie(nullptr);
}

int main()
{
    Init();

    int NumCase = 0;
    std::cin >> NumCase;

    std::vector<int> Answers(NumCase);

    for (int Case = 0; Case < NumCase; Case++)
    {
        int ArraySize = 0;
        std::cin >> ArraySize;

        std::vector<long long> Array(ArraySize);
        for (int i = 0; i < ArraySize; i++)
        {
            std::cin >> Array[i];
        }

        int GoodCount = 0;

        long long MaxNum = 0;
        long long Sum = 0;

        for (int i = 0; i < ArraySize; i++)
        {
            Sum += Array[i];
            MaxNum = std::max(Array[i], MaxNum);

            if (MaxNum * 2 == Sum)
            {
                GoodCount++;
            }
        }

        Answers[Case] = GoodCount;
    }

    for (int i = 0; i < NumCase; i++)
    {
        std::cout << Answers[i] << "\n";
    }

    return 0;
}