문제

Programmers, 주식가격

핵심

• 입력의 크기가 $10^5$이라 대략 $O(N^2)$ 미만 알고리즘을 사용한다.
• stack으로 풀고 싶어서 여러 방법을 고민해 봤지만, 방법이 떠오르지 않았다. 시간초과라도 맛보자는 생각에 $O(N^2)$으로 편하게 구현했는데 이게 왠 걸? 맞는다. 적어도 효율성 케이스에선 시간초과가 나도록 해야 하지 않았을까? level2 문제라 자유로운 풀이를 허용하나보다,,
• 이 문제에서 배워야 하는 건 스택으로 푸는 $O(N)$ 풀이다. 스택에는 해당 원소의 인덱스를 저장하다가 가격이 내려가는 순간 스택의 원소를 pop하여 가격이 내려가지 않았던 기간을 구한다. 현재 인덱스와 스택의 인덱스를 비교하여 내려가지 않았던 기간을 구할 수 있다. 코드로 표현하면 아래와 같다.

// prices 배열 순회
while (!s.isEmpty() && prices[i] < prices[s.peek()])
	answer[s.peek()] = i - s.pop();   

• 가격이 내려간 원소에 대한 가격이 내려지지 않은 기간을 위에서 계산한 후, 스택에 남아 있는 원소에 대해 가격이 내려가지 않은 기간을 구한다. 처음부터 끝까지 가격이 내려가지 않았다면 전체 배열의 길이가 되므로, 전체 길이에서 해당 인덱스 + 1만큼 차감해 준다.

while (!s.isEmpty())
	answer[s.peek()] = len - 1 - s.pop();

개선

#include <string>
#include <vector>
#include <stack>
#include <iostream>
using namespace std;

vector<int> solution(vector<int> prices) {
    vector<int> answer(prices.size(), 0);
    for (int i = 0; i < prices.size(); ++i) {
        int period = 0;
        for (int j = i + 1; j < prices.size(); ++j) {
            ++period;
            if (prices[i] > prices[j])
                break;
        }
        answer[i] = period;
    }
    answer[prices.size() - 1] = 0;
    return answer;
}

코드

시간복잡도

• $O(N)$

C++

#include <string>
#include <vector>
#include <stack>
#include <iostream>

using namespace std;
vector<int> solution(vector<int> prices) {
    int len = prices.size();
    vector<int> answer(len, 0);
    stack<int> s;
    for (int i = 0; i < len; ++i) {
        while (!s.empty() && prices[i] < prices[s.top()]) {
            answer[s.top()] = i - s.top();
            s.pop();
        }
        s.push(i);
    }
    while (!s.empty()) {
        answer[s.top()] = len - 1 - s.top();
        s.pop();
    }
    return answer;
}

Java

import java.util.*;

class Solution {
    public int[] solution(int[] prices) {
        int len = prices.length;
        int[] answer = new int[len];
        Stack<Integer> s = new Stack<>();
        for (int i = 0; i < len; ++i) {
            while (!s.isEmpty() && prices[i] < prices[s.peek()])
                answer[s.peek()] = i - s.pop();   
            s.push(i);
        }
        while (!s.isEmpty())
            answer[s.peek()] = len - 1 - s.pop();
        return answer;
    }
}