[프로그래머스] 주식가격

문제 링크 : https://school.programmers.co.kr/learn/courses/30/lessons/42584

접근 방식

이중 for문을 통해, 시작점과 끝점을 정해 어디까지 현재 가격을 갈 수 있는지를 확인한다.

이후, answer 배열에 끝점-시작점을 담는다.

결과

소스 코드

import java.util.*;
class Solution {
    public int[] solution(int[] prices) {
        int[] answer = new int[prices.length];
        for(int i=0; i<prices.length; i++){
            int cur = prices[i];
            int start = i;
            int end = i;
            for(int j=start+1; j<prices.length; j++){
                if(cur<=prices[j]) end = j;
                else {
                    end=j;
                    break;
                }
            }
            answer[i] = end-start;
        }

        return answer;
    }
}

결과는 통과가 나왔다. 사실 현재 코드는 시간 복잡도가 O(n^2)로 prices의 길이가 최대 10만이기에 무조건 시간 초과가 발생해야하는데, 테스트케이스가 극한의 조건을 사용하지 않아 통과과 된 걸로 보인다.

다른 접근 방식

최적화

물론 문제는 통과가 됐지만, 잘못 만들어진 문제라고 판단이 되어서 코드를 최적화할 수 있어야 한다고 생각했다.

두 번째 start와 end를 찾는 방식을 for문을 통한 순회가 아닌, 다른 방식을 사용해야 한다.

프로그래머스의 뒤에 있는 큰 수 찾기 (블로그 문제 풀이) 문제와 비슷한 풀이로 접근할 수 있다.
스택 자료구조를 이용해서 인덱스를 관리해서, 스택에서 꺼낼 수 인덱스에 해당하는 값이 현재의 값보다 큰지 여부를 통해 확인할 수 있다.

import java.util.*;

class Solution {
    public int[] solution(int[] prices) {
        int[] answer = new int[prices.length];
        Stack<Integer> stack = new Stack<>();

        // 주식 가격을 하나씩 확인
        for (int i = 0; i < prices.length; i++) {
            // 주식이 떨어지는 경우, 바로 계산
            while (!stack.isEmpty() && prices[i] < prices[stack.peek()]) {
                int idx = stack.pop();
                answer[idx] = i - idx; // 가격이 떨어지기까지의 시간 계산
            }
            stack.push(i); // 현재 시점을 스택에 저장
        }

        // 현재 stack에 있는 인덱스는 주식 가격이 떨어지지 않은 인덱스
        while (!stack.isEmpty()) {
            int idx = stack.pop();
            answer[idx] = prices.length - idx - 1;
        }

        return answer;
    }
}

크게 두 번의 로직을 통해 계산을 진행한다.

첫번째 while문의 경우, 직전 가격과 현재 가격을 비교해서 만약 현재 가격이 더 작다면 (주식 하락) pop을 통해 해당 인덱스를 반환한다. 이후에는 해당 인덱스의 시간을 계산해서 저장한다.

두번째 while문은 시간이 종료될 때까지 주식이 하락하지 않은 경우를 계산하기 위함이다. 주식이 하락하지 않은 경우 담겨있는 길이-인덱스-1를 통해 유지 시간을 확인할 수 있다.

이 로직의 핵심은 stack 자료구조를 이용해서 인덱스(시간)를 관리할 줄 알아야한다는 것이다. stack 자료구조의 특성상 직전 값과의 비교과 매우 수월하고 빠르게 진행되기 때문에 이 문제에서는 해당 초에서의 가격이 얼마나 유지되는지를 파악하는 것이 아니라, 언제 끊기는지를 확인하고 끊기지 않은 경우를 위한 로직을 설계하면 쉽게 풀 수 있다.