Java O(n * log(k)) solution, with a wrapped up TreeMap


  • 2
    Z
    • Wrap up TreeMap into TreeMapDup to
      • hold duplicates
      • insertion/deletion/query min or max in log(K)
      • get total size in O(1)
    • hold the data inside the window in 2 TreeMapDup left and right, maintain left.size >= right.size
    • when data moving out and moving in, rebalance left and right to maintain left.size >= right.size
        public double[] medianSlidingWindow(int[] nums, int k) {
            int n = nums.length;
            double[] result = new double[n - k + 1];
    
            TreeMapDup left = new TreeMapDup();
            TreeMapDup right = new TreeMapDup();
    
            for (int i = 0; i < n; i++) {
                left.add(nums[i]);
                int leftMax = left.lastKey();
                left.remove(leftMax);
                right.add(leftMax);
                reBalance(left, right);
    
                int l = i - k + 1;
                if (l < 0) {
                    continue;
                }
    
                result[l] = (double)left.lastKey();
                if (k % 2 == 0) {
                    result[l] = (result[l] + (double)right.firstKey()) / 2.0;
                }
                int old = nums[l];
                if (!left.remove(old)) {
                    right.remove(old);
                }
                reBalance(left, right);
            }
    
            return result;
        }
    
        private void reBalance(TreeMapDup left, TreeMapDup right) {
            if (left.size < right.size) {
                int rightMin = right.firstKey();
                right.remove(rightMin);
                left.add(rightMin);
            }
        }
    
        class TreeMapDup {
            int size;
            private final TreeMap<Integer, Integer> treeMap;
            public TreeMapDup() {
                treeMap = new TreeMap<>();
                size = 0;
            }
    
            int lastKey() {
                return treeMap.lastKey();
            }
    
            int firstKey() {
                return treeMap.firstKey();
            }
    
            void add(int key) {
                size++;
                treeMap.put(key, treeMap.getOrDefault(key, 0) + 1);
            }
    
            boolean remove(int key) {
                if (!treeMap.containsKey(key)) {
                    return false;
                }
    
                size--;
                int size = treeMap.get(key);
                if (size == 1) {
                    treeMap.remove(key);
                } else {
                    treeMap.put(key, size - 1);
                }
    
                return true;
            }
        }
    

Log in to reply
 

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