[알고리즘 - 그리디] 주식

문제

문제 링크 - https://www.acmicpc.net/problem/11501

접근법

정답을 얻기 위해서는 아래와 같이 계산하면 된다.

1. 현재 주식 가격 이후에 더 비싼 가격이 있으면 -> 오늘 주식을 산다
2. 현재 주식이 남은 주식 가격 중 가장 비싼 가격이면 -> 현재까지 구매한 주식을 모두 판다
3. 앞으로 더 비싼 주식 가격이 없으면 -> 아무것도 하지 않는다

문제는 이를 구현하는 방법인데, 처음에는 단순히 (주식 가격) 배열의 앞에서부터 순회하며 계산을 시도했다. 코드는 아래처럼 될 것이다.

// main()에서 입력 처리 후 solution을 부르도록 작성
public static int solution(final int days, final int[] prices) {
    PriorityQueue<Integer> queue = new PriorityQueue<>(Comparator.reverseOrder());
    for (int price : prices) {
        queue.add(price);
    }

    long answer = 0;
    List<Integer> currentStocks = new ArrayList<>();
    for (int price : prices) {
        if (!queue.isEmpty() && price < queue.peek()) { // 이후에 더 비싼 주식 존재
            // 주식 구매 로직
            currentStocks.add(price); 
        } else if (!queue.isEmpty() && price == queue.peek()) { // 남은 주식 중 최고가인 경우
            for (Integer integer : currentStocks) {
                // 주식 판매 로직
                answer += (price - integer);
            }
            // 구매했던 주식들 판매 완료했으므로 clear
            currentStocks.clear();
        }
        // 주식 가격에 관계없이 순회가 끝난 값은 우선순위큐에서 뺴주어야 함
        queue.remove(price);
    }
    return answer;
}

남은 주식 중 가장 비싼 가격을 찾기 위해 우선순위 큐에 주식 가격들을 역순으로 넣고, queue.peek() 으로 남은 주식 가격 중 최고가를 조회하면서 현재 주식 가격이 더 작으면 구매한다(list에 add). 그러다가 남은 주식 중 최고가를 만나면 지금까지 구매했던 주식들을 모두 판매한다.

이렇게 계산하면 주어진 테스트 케이스에 올바른 답이 나오긴하나, 시간초과로 문제가 풀리지 않는다. 우선순위큐를 생성하는 비용과 매번 제거하는 비용, 주식을 판매할 때 순회하는 비용이 그 원인이었다.

배열의 앞에서부터 순회하면 항상 위와 같이 비효율적인 (구매한 주식을 팔기 위한 순회 등) 로직이 등장할 수밖에 없어보였다. 이 문제를 해결하기 위해서는 배열의 뒤에서부터 순회하면 된다. 뒤에서부터 순회하게 되면 주식을 팔아야 할 시점과 주식을 판매할 때의 주식 가격을 미리 알 수 있으므로 배열을 한 번 순회하기만 하면 최종 이익을 계산할 수 있으므로 O(N)에 문제를 해결할 수 있다.

그림으로 표현하면 아래와 같다.

앞에서 부터 순회

(이후 생략)

때문에 주식을 판매할 때 순회가 발생하고 매번 남은 주식 중 최고가를 계산해야 한다

뒤에서 부터 순회

(이후 생략)

때문에 뒤에서부터 한번만 순회하면 전체에서 최고 판매 이득을 구할 수 있게 된다

위 로직을 코드로 작성하면 아래와 같이 된다.

public static long solution(final int days, final int[] prices) {
    long count = 0;
    int prevMax = Integer.MIN_VALUE;

    for (int i = prices.length - 1; i >= 0; i--) {
        int current = prices[i];
        if (current > prevMax) {
            prevMax = current;
            continue;
        }

        if (current < prevMax) {
            count += (prevMax - current);
        }
    }

    return count;
}

한가지 주의할점은 입력 값의 범위가 2 <= N <= 1,000,000 이고 0 <= 주가 <= 10,000 이므로 이익이 int 범위를 벗어날 수 있어서 count를 int로 선언할 경우 오답 처리된다. (문제에서도 답이 부호있는 64bit 정수형으로 표현 가능하다고 언급되어 있음) 따라서 count의 타입을 long으로 선언해주어야 문제 없이 정답 처리된다.