[Java/Python] 백준 11501 주식

문제

11501번: 주식 (acmicpc.net)

11501번: 주식

풀이

우선순위 큐를 이용하여 풀이한 로직은 다음과 같다.  O ( N^2)

1. 우선순위 큐에 배열값을 다 담는다.
2. 방문 해시맵을 만들어 아직 남아있는 개수이면, 해당 key 방문처리하고 꺼낸다. 
3. 현재 값이 max Number 보다 작으면 합을 더한다. 
4. 현재 인덱스가 maxIdx보다 크면, 최댓값을 for문으로 구한다. 
5. 배열 끝까지 2..4 로직을 반복하고 답을 구한다. 

그런데 이 풀이법은 O(N^2)라서 주어진 배열 길이 백만에서는 어림도 없다. 

 

O(N) 풀이법이 있다고?  다음과 같다. 

우리가 주식을 보면 표를 정방향으로도 보고 거꾸로 뒤집어도 본다. 거꾸로 탐색하는 방법을 생각해보자 .

 

로직은 간단하다. 리버스 탐색에서 현재 max N 보다 작으면 차이 만큼 덧셈한다. 그게 아니라면 최댓값 갱신. 

하지만 정방향으로만 보는 고정관념을 깨는 풀이법이어서 힌트를 볼때까지 떠오르지 못했다.

 

참고

구현 코드도 중요하니(사실 짜 논게 아까워서) 첨부한다. 

코드

파이썬

class Main():
    def __init__(self):
        for _ in range(int(input())):
            n= int(input())
            nums = list(map(int, input().split()))
            print(self.getN(n,nums))

    def getN(self,n,nums):
        ret = 0
        maxN = nums[-1]
        for i in range(n-2, -1, -1):
            v = nums[i]
            if v < maxN : ret += (maxN - v)
            else : maxN = v
        return ret

a = Main()

 

자바

O(N^2) 풀이 

import java.io.*;
import java.security.cert.CollectionCertStoreParameters;
import java.util.*;

public class Main {

    public static void main(String[] args) throws IOException {
        Main m = new Main();
        m.solution();

    }

    public long helper(int[] nums) {
        long ret = 0;
        //{idx, value}
        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> b[1] - a[1]==0 ? a[0]- b[0] : b[1]-a[1]);
        HashMap<Integer, Integer> map = new HashMap<>();
        for (int i = 0; i < nums.length; i++) {
            map.put(nums[i], map.getOrDefault(nums[i], 0) + 1);
            pq.offer(new int[]{i, nums[i]});
        }
        int len = nums.length;
        int info[] = pq.poll();
        int maxN = info[1];
        int maxIdx = info[0];

        map.put(maxN, map.get(maxN) - 1);
        for (int i = 0; i < len-1; i++) {
            if (map.get(nums[i]) != 0) {
                if (i >= maxIdx) {
                    int a = 0;
                    int l = nums.length;
                    while (!pq.isEmpty()) {
                        int temp[] = pq.poll();
                        if (map.get(temp[1]) > 0 && temp[0] > i) {
                            l = temp[0];
                            a = temp[1];
                            map.put(a, map.get(a) - 1);
                            maxN = a;
                            maxIdx = l;
                            break;
                        }
                    }
                    if(l == nums.length)return ret;
                }
                map.put(nums[i], map.get(nums[i]) - 1);
                if (nums[i] < maxN) {
                    ret += maxN - nums[i];
                }
            }
            //System.out.println(map);
            //System.out.println("i"+i+" ret:"+ret+"info[]:"+maxIdx+","+maxN);
        }
        return ret;
    }

    public void solution() throws IOException {
        FastReader fr = new FastReader();
        int T = fr.nextInt();
        String ans = "";
        while (T-- > 0) {
            int N = fr.nextInt();
            int nums[] = new int[N];
            for (int i = 0; i < N; i++) {
                nums[i] = fr.nextInt();
            }
            ans += String.valueOf(helper(nums));
            ans += "\n";
        }
        System.out.print(ans);
    }

    public static class FastReader {
        BufferedReader br;
        StringTokenizer st;

        public FastReader() {
            br = new BufferedReader(new InputStreamReader(System.in));
        }

        public FastReader(String s) throws FileNotFoundException {
            br = new BufferedReader(new FileReader(new File(s)));
        }

        String next() {
            while (st == null || !st.hasMoreElements()) {
                try {
                    st = new StringTokenizer(br.readLine());
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            return st.nextToken();
        }

        int nextInt() {
            return Integer.parseInt(next());
        }

        long nextLong() {
            return Long.parseLong(next());
        }

        double nextDouble() {
            return Double.parseDouble(next());
        }

        String nextLine() {
            String str = "";
            try {
                str = br.readLine();
            } catch (IOException e) {
                e.printStackTrace();
            }
            return str;
        }
    }
}

O(N)풀이

import java.io.*;
import java.security.cert.CollectionCertStoreParameters;
import java.util.*;

public class Main {

    public static void main(String[] args) throws IOException {
        Main m = new Main();
        m.solution();
    }

    public long helper(int[] nums) {
        long ret = 0;
        int maxN = nums[nums.length-1];

        for(int i = nums.length-2 ;i>=0 ; i--){
                if(nums[i] < maxN)ret+=(maxN - nums[i]);
                else maxN = nums[i];
        }
        return ret;
    }

    public void solution() throws IOException {
        FastReader fr = new FastReader();
        int T = fr.nextInt();
        String ans = "";
        while (T-- > 0) {
            int N = fr.nextInt();
            int nums[] = new int[N];
            for (int i = 0; i < N; i++) {
                nums[i] = fr.nextInt();
            }
            ans += String.valueOf(helper(nums));
            ans += "\n";
        }
        System.out.print(ans);
    }

    public static class FastReader {
        BufferedReader br;
        StringTokenizer st;

        public FastReader() {
            br = new BufferedReader(new InputStreamReader(System.in));
        }

        public FastReader(String s) throws FileNotFoundException {
            br = new BufferedReader(new FileReader(new File(s)));
        }

        String next() {
            while (st == null || !st.hasMoreElements()) {
                try {
                    st = new StringTokenizer(br.readLine());
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            return st.nextToken();
        }

        int nextInt() {
            return Integer.parseInt(next());
        }

        long nextLong() {
            return Long.parseLong(next());
        }

        double nextDouble() {
            return Double.parseDouble(next());
        }

        String nextLine() {
            String str = "";
            try {
                str = br.readLine();
            } catch (IOException e) {
                e.printStackTrace();
            }
            return str;
        }
    }
}