fd

Author: dingjikerbo

solution/src/main/java/com/inuker/solution/SlidingWindowMedian.java

@@ -1,80 +1,88 @@
 package com.inuker.solution;
 
 import java.util.Collections;
+import java.util.Comparator;
+import java.util.Map;
 import java.util.TreeMap;
 
 /**
  * Created by liwentian on 2017/9/12.
+ * https://leetcode.com/articles/sliding-window-median/
+ * 和 #295 Find Median from Data Stream 比较类似
  */
 
 public class SlidingWindowMedian {
 
+    /**
+     * 这题如果用PriorityQueue则复杂度为O(nk)，因为其删除复杂度为O(n)
+     * 改为TreeMap会降为O(nlgk)
+     */
     public double[] medianSlidingWindow(int[] nums, int k) {
-        double[] res = new double[nums.length - k + 1];
-        TreeMap<Integer, Integer> minHeap = new TreeMap<Integer, Integer>();
-        TreeMap<Integer, Integer> maxHeap = new TreeMap<Integer, Integer>(Collections.reverseOrder());
+        MyMap up = new MyMap();
+        MyMap down = new MyMap(Comparator.reverseOrder());
+        double[] result = new double[nums.length - k + 1];
+        for (int i = 0; i < nums.length; i++) {
+            up.myAdd(nums[i]);
+            down.myAdd(up.myPollFirst());
 
-        int minHeapCap = k / 2; //smaller heap when k is odd.
-        int maxHeapCap = k - minHeapCap;
+            if (i >= k) {
+                if (up.containsKey(nums[i - k])) {
+                    up.myRemove(nums[i - k]);
+                } else {
+                    down.myRemove(nums[i - k]);
+                }
+            }
 
-        for (int i = 0; i < k; i++) {
-            maxHeap.put(nums[i], maxHeap.getOrDefault(nums[i], 0) + 1);
+            if (up.size < down.size) {
+                up.myAdd(down.myPollFirst());
+            }
+            if (i >= k - 1) {
+                result[i - k + 1] = up.size == down.size
+                        ? ((double) up.firstKey() + down.firstKey()) / 2 : up.firstKey();
+            }
         }
-        int[] minHeapSize = new int[]{0};
-        int[] maxHeapSize = new int[]{k};
-        for (int i = 0; i < minHeapCap; i++) {
-            move1Over(maxHeap, minHeap, maxHeapSize, minHeapSize);
+        return result;
+    }
+
+    class MyMap extends TreeMap<Integer, Integer> {
+        int size;
+
+        MyMap() {
+            super();
         }
 
-        res[0] = getMedian(maxHeap, minHeap, maxHeapSize, minHeapSize);
-        int resIdx = 1;
+        MyMap(Comparator<Integer> comparator) {
+            super(comparator);
+        }
 
-        for (int i = 0; i < nums.length - k; i++) {
-            int addee = nums[i + k];
-            if (addee <= maxHeap.keySet().iterator().next()) {
-                add(addee, maxHeap, maxHeapSize);
-            } else {
-                add(addee, minHeap, minHeapSize);
-            }
+        void myAdd(int n) {
+            put(n, getOrDefault(n, 0) + 1);
+            size++;
+        }
 
-            int removee = nums[i];
-            if (removee <= maxHeap.keySet().iterator().next()) {
-                remove(removee, maxHeap, maxHeapSize);
+        boolean myRemove(int n) {
+            int count = getOrDefault(n, 0);
+            if (count == 0) {
+                return false;
+            }
+            if (count == 1) {
+                remove(n);
             } else {
-                remove(removee, minHeap, minHeapSize);
+                put(n, count - 1);
             }
+            size--;
+            return true;
+        }
 
-            //rebalance
-            if (minHeapSize[0] > minHeapCap) {
-                move1Over(minHeap, maxHeap, minHeapSize, maxHeapSize);
-            } else if (minHeapSize[0] < minHeapCap) {
-                move1Over(maxHeap, minHeap, maxHeapSize, minHeapSize);
+        int myPollFirst() {
+            Map.Entry<Integer, Integer> entry = firstEntry();
+            if (entry.getValue() == 1) {
+                pollFirstEntry();
+            } else {
+                put(entry.getKey(), entry.getValue() - 1);
             }
-
-            res[resIdx] = getMedian(maxHeap, minHeap, maxHeapSize, minHeapSize);
-            resIdx++;
+            size--;
+            return entry.getKey();
         }
-        return res;
-    }
-
-    public double getMedian(TreeMap<Integer, Integer> bigHeap, TreeMap<Integer, Integer> smallHeap, int[] bigHeapSize, int[] smallHeapSize) {
-        return bigHeapSize[0] > smallHeapSize[0] ? (double) bigHeap.keySet().iterator().next() : ((double) bigHeap.keySet().iterator().next() + (double) smallHeap.keySet().iterator().next()) / 2.0;
-    }
-
-    //move the top element of heap1 to heap2
-    public void move1Over(TreeMap<Integer, Integer> heap1, TreeMap<Integer, Integer> heap2, int[] heap1Size, int[] heap2Size) {
-        int peek = heap1.keySet().iterator().next();
-        add(peek, heap2, heap2Size);
-        remove(peek, heap1, heap1Size);
-    }
-
-    public void add(int val, TreeMap<Integer, Integer> heap, int[] heapSize) {
-        heap.put(val, heap.getOrDefault(val, 0) + 1);
-        heapSize[0]++;
-    }
-
-    public void remove(int val, TreeMap<Integer, Integer> heap, int[] heapSize) {
-        if (heap.put(val, heap.get(val) - 1) == 1) heap.remove(val);
-        heapSize[0]--;
     }
 }

test/src/main/java/com/inuker/test/Test2.java

@@ -0,0 +1,81 @@
+package com.inuker.test;
+
+import java.util.Comparator;
+import java.util.Map;
+import java.util.TreeMap;
+
+/**
+ * Created by liwentian on 2017/12/5.
+ */
+
+public class Test2 {
+
+    public double[] medianSlidingWindow(int[] nums, int k) {
+        MyMap up = new MyMap();
+        MyMap down = new MyMap(Comparator.reverseOrder());
+        double[] result = new double[nums.length - k + 1];
+        for (int i = 0; i < nums.length; i++) {
+            up.myAdd(nums[i]);
+            down.myAdd(up.myPollFirst());
+
+            if (i >= k) {
+                if (up.containsKey(nums[i - k])) {
+                    up.myRemove(nums[i - k]);
+                } else {
+                    down.myRemove(nums[i - k]);
+                }
+            }
+
+            if (up.size < down.size) {
+                up.myAdd(down.myPollFirst());
+            }
+            if (i >= k - 1) {
+                result[i - k + 1] = up.size == down.size
+                        ? ((double) up.firstKey() + down.firstKey()) / 2 : up.firstKey();
+            }
+        }
+        return result;
+    }
+
+    class MyMap extends TreeMap<Integer, Integer> {
+        int size;
+
+        MyMap() {
+            super();
+        }
+
+        MyMap(Comparator<Integer> comparator) {
+            super(comparator);
+        }
+
+        void myAdd(int n) {
+            put(n, getOrDefault(n, 0) + 1);
+            size++;
+        }
+
+        boolean myRemove(int n) {
+            int count = getOrDefault(n, 0);
+            if (count == 0) {
+                return false;
+            }
+            if (count == 1) {
+                remove(n);
+            } else {
+                put(n, count - 1);
+            }
+            size--;
+            return true;
+        }
+
+        int myPollFirst() {
+            Map.Entry<Integer, Integer> entry = firstEntry();
+            if (entry.getValue() == 1) {
+                pollFirstEntry();
+            } else {
+                put(entry.getKey(), entry.getValue() - 1);
+            }
+            size--;
+            return entry.getKey();
+        }
+    }
+}

test/src/main/java/com/inuker/test/main.java

@@ -15,18 +15,23 @@
 import java.util.LinkedHashMap;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.Map;
 import java.util.PriorityQueue;
 import java.util.Queue;
 import java.util.Random;
 import java.util.Set;
 import java.util.Stack;
 import java.util.TreeMap;
+import java.util.TreeSet;
 
 public class main {
 
     public static void main(String[] args) {
-        Test test = new Test();
-        test.call("hello");
-        test.call("how");
+        double[] f = new Test2().medianSlidingWindow(new int[] {
+                Integer.MAX_VALUE, 1, 2, 3, 4, 5, 6, 7, Integer.MAX_VALUE
+        }, 2);
+        for (double d : f) {
+            System.out.print(d + " ");
+        }
     }
 }