Day37: Add more test cases and fix assumption about how to deal with tie breakers

Author: thatguyintech

day37/topKFrequent.py

@@ -9,8 +9,7 @@ def topKFrequent(nums, k):
         else:
             freqs[num] = 1
 
-    buckets = [None for i in xrange(max(freqs.keys())+1)]
-
+    buckets = [None for i in xrange(max(freqs.values())+1)]
     for num in freqs:
         freq = freqs[num]
         if buckets[freq]:
@@ -22,21 +21,26 @@ def topKFrequent(nums, k):
     i = len(buckets)-1
     addedCount = 0
     offset = 0
-    while addedCount < k:
-        while not buckets[i-offset-addedCount]:
+    while addedCount < k and i-offset > 0:
+        while not buckets[i-offset]:
             offset += 1
-        ret += buckets[i-offset-addedCount]
-        addedCount += 1
+        ret += buckets[i-offset]
+        addedCount += len(buckets[i-offset])
+        offset += 1
 
     return ret
 
 def testTopKFrequent():
     assert set(topKFrequent([], 0)) == set([])
     assert set(topKFrequent([1], 1)) == set([1])
+    assert set(topKFrequent([-1, -1], 1)) == set([-1])
     assert set(topKFrequent([1,1,1,2,2,3], 2)) == set([1, 2])
     assert set(topKFrequent([-1,-1,-1,2,2,3], 2)) == set([-1, 2])
     assert set(topKFrequent([1,1,1,2,2,3], 3)) == set([1, 2, 3])
+    assert set(topKFrequent([1,1,1,2,2,2,3,3,3], 3)) == set([1, 2, 3])
     assert set(topKFrequent([1, 2, 3, 4, 5], 1)) == set([1, 2, 3, 4, 5])
+    assert set(topKFrequent([4,1,-1,2,-1,2,3], 2)) == set([-1, 2])
+    assert set(topKFrequent([3,2,3,1,2,4,5,5,6,7,7,8,2,3,1,1,1,10,11,5,6,2,4,7,8,5,6], 10)) == set([1,2,5,3,7,6,4,8,10,11])
 
 def main():
     testTopKFrequent()