[골드4] 백준 17298 : 오큰수 - C/C++

문제

크기가 N인 수열 A = A1, A2, ..., AN이 있다. 수열의 각 원소 Ai에 대해서 오큰수 NGE(i)를 구하려고 한다. Ai의 오큰수는 오른쪽에 있으면서 Ai보다 큰 수 중에서 가장 왼쪽에 있는 수를 의미한다. 그러한 수가 없는 경우에 오큰수는 -1이다.

예를 들어, A = [3, 5, 2, 7]인 경우 NGE(1) = 5, NGE(2) = 7, NGE(3) = 7, NGE(4) = -1이다. A = [9, 5, 4, 8]인 경우에는 NGE(1) = -1, NGE(2) = 8, NGE(3) = 8, NGE(4) = -1이다.

 

입력

첫째 줄에 수열 A의 크기 N (1 ≤ N ≤ 1,000,000)이 주어진다. 둘째 줄에 수열 A의 원소 A1, A2, ..., AN (1 ≤ Ai ≤ 1,000,000)이 주어진다.

 

출력

총 N개의 수 NGE(1), NGE(2), ..., NGE(N)을 공백으로 구분해 출력한다.

 

 

 

문제를 처음 접하면 가장 쉽게 생각할 수 있는게 이중 for문을 도는 방식입니다.

배열을 순회하면서 자신의 오른쪽에 있는 수들을 조사하면서 자신보다 큰 수가 나오면 그 수를 오큰수로 판정합니다.

예를 들어, 1 7 2 9 8이 입력으로 들어왔다면 이들을 배열에 저장하고 배열을 순회하면서 1은 7 2 9 8을 검사하고, 7은 2 9 8을 검사하고 2는 9 8을 검사하고 9는 8을 검사합니다. 결과는 각각 1은 7,  7은 9,  2는 9,  9는 오른쪽에 자신보다 더 큰수가 없으므로 -1,  8은 가장 오른쪽에 있으므로 -1입니다.

 

이중 for문을 돌면서 오큰수를 찾는 버전 코드

std::vector<int> _vec(n);
for (int i = 0; i < n; i++)
{
    std::cin >> _vec[i];
}

// 벡터를 순회하면서 i보다 큰 인덱스(오른쪽에 있는)를 가지는
// 벡터의 원소들을 순회하면서 자신보다 큰 수가 있다면
// 그 수가 오큰수, 없다면 -1을 출력한다.
for (int i = 0; i < n; i++)
{
    int rightNum = _vec[i];
    for (int j = i + 1; j < n; j++)
    {
        if (_vec[i] < _vec[j])
        {
            rightNum = _vec[j];
            break;
        }
    }
	
    // 출력
    if (rightNum == _vec[i])
    {
        std::cout << "-1 ";
    }
    else
    {
        std::cout << rightNum << ' ';
    }
}

하지만 이런 방식은 이중 for문으로 배열을 순회하므로 O(n²)의 시간복잡도를 가집니다. 따라서 시간 안에 풀 수 없습니다.

 

 

다음 방식은 스택을 이용하는 방식입니다.

입력으로 1 7 2 9 8이 들어온다고 한다면, 1이 들어왔을 때 스택에 넣습니다. 스택에는 아무 것도 없으므로 비교하지 않습니다. 1을 다음 인풋과 비교하기 위해 prev 변수에 저장합니다.

7이 들어오면, prev인 1과 비교합니다. 이전에 들어온 숫자보다 크다면, 현재 들어온 숫자는 이전 숫자의 오큰수입니다. 또, 스택에 있는 다른 숫자들의 오큰수일 수 있습니다. 따라서 스택을 top부터 하나씩 꺼내면서 현재 들어온 7과 비교해 7이 더 크다면 그 숫자들을 pop하고, 오큰수를 7이라고 기록해놓습니다. 따라서 7은 1의 오큰수가 됩니다. 1을 스택에서 pop합니다. 7 또한 다음에 들어올 인풋들과 비교해야하기 때문에 스택에 push합니다.

2가 들어오면, prev인 7과 비교합니다. 더 작으므로, 2는 7의 오큰수가 아닙니다. 2를 스택에 push합니다.

9가 들어오면, prev인 2와 비교합니다. 더 크므로 9는 2의 오큰수입니다. 2를 스택에서 pop합니다. 또 9는 7보다 크므로, 9는 7의 오큰수입니다. 7을 pop합니다. 9를 다음 인풋들과 비교하기 위해 스택에 push합니다....

(모든 인풋이 들어올 때까지 위 과정을 반복합니다., 가장 마지막 인풋은 무조건 -1)

 

 

스택을 사용한 성공 버전 코드

const int Max = 1000001;

std::stack<std::pair<int, int>> _stack;	// 인풋과 비교할 숫자들, 두 번째는 입력받은 순서
std::vector<int> _resultVec(n);			// 결과를 저장

int prev = Max;
for (int i = 0; i < n; i++)
{
    int input;
    std::cin >> input;

    // 바로 이전 숫자와 현재 숫자 비교해서 현재 숫자가 더 크다면
    if (prev < input)
    {
        // 스택을 돌면서 현재 숫자보다 작은 것들은 전부 빼주고 결과입력
        while (!_stack.empty())
        {
            std::pair<int, int> top = _stack.top();

            if (top.first >= input)
                break;

            _resultVec[top.second] = input;
            _stack.pop();
        }
    }
    else
    {
        prev = input;
    }

    // 현재 숫자를 스택에 넣음 (다음 번에 입력되는 숫자와 체크 필요)
    _stack.push(std::pair<int, int>(input, i));
}

// 결과를 출력
for (const auto& it : _resultVec)
{
    // 결과 값이 없으면 -1
    if (it == 0)
    {
        std::cout << "-1 ";
    }
    else
    {
        std::cout << it << ' ';
    }
}