sofkyle
diff --git a/‎LeetCodePrj/Java/leetcode/hard/page2/HardPage2.java‎
Lines changed: 7 additions & 6 deletions b/‎LeetCodePrj/Java/leetcode/hard/page2/HardPage2.java‎
Lines changed: 7 additions & 6 deletions
diff --git a/‎LeetCodePrj/Java/leetcode/hard/page2/LFUCache.java‎
Lines changed: 38 additions & 10 deletions b/‎LeetCodePrj/Java/leetcode/hard/page2/LFUCache.java‎
Lines changed: 38 additions & 10 deletions
@@ -5,16 +5,17 @@
 public class HardPage2 {
     @Test
     public void testLFUCache() {
-        LFUCache lfuCache = new LFUCache(2);
-        lfuCache.set(1,1);
-        lfuCache.set(2,2);
+        LFUCache lfuCache = new LFUCache(3);
+        lfuCache.put(2,2);
+        lfuCache.put(1,1);
+        lfuCache.get(2);
         lfuCache.get(1);
-        lfuCache.set(3,3);
         lfuCache.get(2);
+        lfuCache.put(3,3);
+        lfuCache.put(4,4);
         lfuCache.get(3);
-        lfuCache.set(4,4);
+        lfuCache.get(2);
         lfuCache.get(1);
-        lfuCache.get(3);
         lfuCache.get(4);
     }
 }
@@ -8,41 +8,69 @@ public class LFUCache {
     // <Key, Counts>
     HashMap<Integer, Integer> counts;
     // <Counts, Key LinkedList>
-    HashMap<Integer, LinkedHashSet<Integer>> lists;
+    HashMap<Integer, LinkedHashSet<Integer>> order;
     int capacity;
     int min = 1;
 
     public LFUCache(int capacity) {
         this.capacity = capacity;
         vals = new HashMap<>();
         counts = new HashMap<>();
-        lists = new HashMap<>();
-        lists.put(1, new LinkedHashSet<>());
+        order = new HashMap<>();
+        order.put(1, new LinkedHashSet<>());
     }
 
     public int get(int key) {
         if (!vals.containsKey(key)) return -1;
 
         int oldCount = counts.get(key);
         counts.put(key, oldCount + 1);
-        if (oldCount == min && lists.get(oldCount).size() == 1) {
+        order.get(oldCount).remove(key);
+        if (oldCount == min && order.get(oldCount).size() == 0) {
             min = oldCount + 1;
         }
-        lists.get(oldCount + 1).add(key);
+        if (!order.containsKey(oldCount + 1)) {
+            order.put(oldCount + 1, new LinkedHashSet<>());
+        }
+        order.get(oldCount + 1).add(key);
 
         return vals.get(key);
     }
 
-    public void set(int key, int value) {
+    public void put(int key, int value) {
+        if (capacity == 0) return;
         if (vals.containsKey(key) || vals.size() < capacity) {
             vals.put(key, value);
-            counts.put(key, counts.get(key) + 1);
-            lists.get(counts.get(key)).add(key);
+            int oldCount = counts.getOrDefault(key, 0);
+            counts.put(key, oldCount + 1);
+
+            if (!order.containsKey(oldCount + 1)) {
+                order.put(oldCount + 1, new LinkedHashSet<>());
+            }
+            order.get(oldCount + 1).add(key);
+            if (oldCount == 0) {
+                min = 1;
+            } else {
+                order.get(oldCount).remove(key);
+                if (oldCount == min && order.get(oldCount).size() == 0) {
+                    min = oldCount + 1;
+                }
+            }
         } else {
-            int removeKey = lists.get(min).iterator().next();
+            // evict
+            int removeKey = order.get(min).iterator().next();
             vals.remove(removeKey);
             counts.remove(removeKey);
-            lists.get(min).remove(lists.get(min).iterator().next());
+            order.get(min).remove(removeKey);
+
+
+            min = 1;
+            vals.put(key, value);
+            counts.put(key, min);
+            if (!order.containsKey(min)) {
+                order.put(min, new LinkedHashSet<>());
+            }
+            order.get(min).add(key);
         }
     }
 }