From O(NKlogK) using heap to O(NK) with explanation


  • 2

    At first, I try to extend Problem 1 directly. In Problem 1, we only have 3 colors (3 states). Each time we need to find minimum value except current color. Now we have K colors, so obviously heap will be helpful. Thus, my first version looks like this.

        // O(N*KlogK) time + O(N) space
        // dp[i][j] = Math.min(any k!= j | dp[i-1][k]) + costs[i][j]
        public int minCostII(int[][] costs) {
            if (costs.length == 0 || costs[0].length == 0) return 0;
            PriorityQueue<Integer> heap = new PriorityQueue<>();
            int[] colors = new int[costs[0].length];
            System.arraycopy(costs[0], 0, colors, 0, costs[0].length);
            
            for (int i = 1; i < costs.length; i++) {
                for (int color : colors) heap.offer(color);
    
                int[] next = new int[colors.length];
                for (int j = 0; j < next.length; j++) { // O(KlogK)
                    heap.remove(colors[j]);
                    next[j] = heap.peek() + costs[i][j];
                    heap.offer(colors[j]);
                }
                colors = next;
                heap.clear();
            }
            
            int min = colors[0];
            for (int i = 1; i < colors.length; i++)
                min = Math.min(min, colors[i]);
            return min;
        }
    

    After reading other O(NK) solutions, I realize heap is too heavy for this problem. We only need to know minimum and second min. Why? Think of a heap [3,4,5,8,9,7,11]. For every color except "3", the min is 3. For "3", min is 4. So the heap remove/offer() above actually operates the bottom of heap most time which is unnecessary absolutely.

    Now it's much more clean and efficient~

        public int minCostII(int[][] costs) {
            if (costs.length == 0 || costs[0].length == 0) return 0;
            int min1 = 0, min2 = 0, idx = -1;
            for (int i = 0; i < costs.length; i++) {
                int min1i = Integer.MAX_VALUE, min2i = min1i, idxi = 0;
                for (int j = 0; j < costs[i].length; j++) {
                    int cost = costs[i][j] + (j != idx ? min1 : min2);
                    if (cost < min1i) {
                        min2i = min1i; min1i = cost; idxi = j;
                    } else if (cost < min2i)
                        min2i = cost;
                }
                min1 = min1i; min2 = min2i; idx = idxi;
            }
            return min1;
        }
    

  • 0
    B

    Thank you for your explanation! I really like the way you describe the heap solution first an then give the optimal solution with two minimum numbers.


  • 0

    @bshi Thanks! Thinking process is more important and interesting than memorizing the optimal solution. Hope all of us can learn something and get fun from this! :)


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.