9/3 to 9/4: Heaps problems

- k_closest_stars.py
- k_largest_in_heap.py
- online_median.py
- sort_almost_sorted_array.py
- sort_increasing_decreasing_array.py
- sorted_arrays_merge.py

Author: jsgoller1

elements-of-programming-interviews/python/k_closest_stars.py

@@ -1,3 +1,5 @@
+import heapq
+import collections
 import functools
 import math
 from typing import Iterator, List
@@ -28,8 +30,17 @@ def __eq__(self, rhs):
 
 
 def find_closest_k_stars(stars: Iterator[Star], k: int) -> List[Star]:
-    # TODO - you fill in here.
-    return []
+    """
+    Could we just create heap, push stars in based on distance to earth, and evict if we store more than k? 
+    O(n * log(k)) time, O(k) space. 
+    """
+
+    heap = []
+    for star in stars:
+        heapq.heappush(heap, (-star.distance, star))
+        if len(heap) > k:
+            heapq.heappop(heap)
+    return [heapq.heappop(heap)[1] for i in range(k)][::-1]
 
 
 def comp(expected_output, output):

elements-of-programming-interviews/python/k_largest_in_heap.py

@@ -2,10 +2,36 @@
 
 from test_framework import generic_test, test_utils
 
+"""
+- In max heap, root is largest, and they only get smaller. 
+- Leaves of item at i are at 2**i, 2**i+1
+- If one child is greater than the other, we can't explore the lesser until we explore the best ones
+- Could we use a heap for this? 
+    - each time we evaluate a node, put it into set of largest, and heappush all children.
+    - heappop to get next node. 
+    - Do this k many times; always ensures that of children, we get the biggest. 
+"""
+
+from collections import namedtuple
+from heapq import heappush, heappop
+
+entry = namedtuple("entry", ["val", "idx"])
+
 
 def k_largest_in_binary_heap(A: List[int], k: int) -> List[int]:
-    # TODO - you fill in here.
-    return []
+    if (not A) or (len(A) <= k):
+        return A
+    largest = []
+    heap = [entry(-A[0], 0)]
+    while len(largest) < k:
+        curr = heappop(heap)
+        largest.append(-curr.val)
+        left_idx, right_idx = 2*curr.idx+1, (2*curr.idx)+2
+        if left_idx < len(A):
+            heappush(heap, entry(-A[left_idx], left_idx))
+        if right_idx < len(A):
+            heappush(heap, entry(-A[right_idx], right_idx))
+    return largest
 
 
 if __name__ == '__main__':

elements-of-programming-interviews/python/online_median.py

@@ -2,10 +2,66 @@
 
 from test_framework import generic_test
 
+"""
+1,0,3,5,2,0,1
+1,0.5,1,2,2,1.5,1
 
-def online_median(sequence: Iterator[int]) -> List[float]:
-    # TODO - you fill in here.
-    return []
+1 -> 1
+0,| 1 -> .5
+0,|1|,3 -> 1
+0,1,| 3,5 -> 2
+0,1,|2|,3,5 -> 2
+0,0,1, | 2,3,5 -> 1.5
+0,0,1,|1|,2,3,5 -> 1
+
+if odd, middle element
+if even, average of two innermost elements. 
+
+brute force: insert, sort, calculate (nlogn)
+
+if we kept the "smallest of the larger half" and "largest of the smaller half" elements, 
+then we can just average them 
+    - can probably look at size of greatests and leasts to determine if odd number or not 
+
+
+"""
+from heapq import heappush, heappop
+
+
+def online_median(arr: Iterator[int]) -> List[float]:
+    medians = []
+    try:
+        first = next(arr)
+        medians.append(first)
+        second = next(arr)
+
+        left = [-first if first < second else -second]
+        right = [first if first > second else second]
+        medians.append((-left[0] + right[0])/2)
+    except StopIteration:
+        return medians
+
+    for val in arr:
+        # ensure every element in left is less than smallest
+        # in right
+        if val < right[0]:
+            heappush(left, -val)
+        else:
+            heappush(right, val)
+
+        # Ensure balance via shifting elements from one
+        # heap to another (in order) if needed
+        while len(left) - len(right) > 1:
+            heappush(right, -heappop(left))
+        while len(right) - len(left) > 1:
+            heappush(left, -heappop(right))
+
+        # Compute median depending on even/odd number of elements
+        if len(right) == len(left):
+            medians.append((-left[0] + right[0])/2)
+        else:
+            medians.append(-left[0] if len(left) > len(right) else right[0])
+    return medians
 
 
 def online_median_wrapper(sequence):

elements-of-programming-interviews/python/sort_almost_sorted_array.py

@@ -2,18 +2,55 @@
 
 from test_framework import generic_test
 
+"""
+ex: k=17, [1, -11, 2, -2, -19, 4, -9, 1, 10, -12, 6, -19, 9, -5, 0, 5, -4, 13, 19, 19, 11]
+return [-19, -19, -12, -11, -9, -5, -4, -2, 0, 1, 1, 2, 4, 5, 6, 9, 10, 11, 13, 19, 19]
+
+- brute force is just sorting the array
+- if k is 0, array is already sorted
+- if k is 1, then between arr[i-1], arr[i], and arr[i+1] there is a minimal element. 
+- the smallest element of the array is within k of 0. 
+- the second smallest element is within k of 1. 
+- if k is greater than or equal to array len, no better than brute force 
+
+"""
+from heapq import heappush, heappop
+
 
 def sort_approximately_sorted_array(sequence: Iterator[int],
                                     k: int) -> List[int]:
-    # TODO - you fill in here.
-    return []
+    heap = []
+    ans = []
+    i = 0
+    while i <= k:
+        try:
+            heappush(heap, next(sequence))
+        except StopIteration:
+            pass
+        i += 1
+    while heap:
+        ans.append(heappop(heap))
+        try:
+            heappush(heap, next(sequence))
+        except StopIteration:
+            continue
+    return ans
 
 
 def sort_approximately_sorted_array_wrapper(sequence, k):
     return sort_approximately_sorted_array(iter(sequence), k)
 
 
 if __name__ == '__main__':
+    cases = [
+        ([1, -11, 2, -2, -19, 4, -9, 1, 10, -12, 6, -19, 9, -5, 0, 5, -4, 13, 19, 19, 11], 17),
+        ([1, 2, 3, 4, 5, 6, 7], 0),
+        ([2, 1, 4, 3, 7, 6, 9, 8], 1),
+    ]
+    for case, k in cases:
+        expected = sorted(case)
+        actual = sort_approximately_sorted_array_wrapper(case, k)
+        assert actual == expected, f"\n{k}:{case}\n{actual} != {expected}"
     exit(
         generic_test.generic_test_main(
             'sort_almost_sorted_array.py', 'sort_almost_sorted_array.tsv',

elements-of-programming-interviews/python/sort_increasing_decreasing_array.py

@@ -2,13 +2,71 @@
 
 from test_framework import generic_test
 
+"""
+Smallest case would be like:
+k = 3: 0 1 0 | 1 0 | 1 0 
 
-def sort_k_increasing_decreasing_array(A: List[int]) -> List[int]:
-    # TODO - you fill in here.
-    return []
+brute force sort is nlogn, can ideally do better 
+vals could be pos, neg, or zero 
+
+"""
+import heapq
+import collections
+
+subarray = collections.namedtuple("subarray", ["val", "left", "right", "increasing"])
+
+
+def get_subarrays(arr):
+    if not arr or (len(arr) == 1):
+        return [arr]
+    subarrays = []
+    curr = []
+    increasing = arr[0] < arr[1]
+    i = 0
+    while i < len(arr)-1:
+        if ((arr[i] < arr[i+1]) and not increasing) or ((arr[i] > arr[i+1]) and increasing):
+            subarrays.append(curr if increasing else curr[::-1])
+            curr = []
+            increasing = not increasing
+        curr.append(arr[i])
+        i += 1
+
+    if not curr:
+        subarrays.append([arr[i]])
+    else:
+        curr.append(arr[i])
+        subarrays.append(curr if increasing else curr[::-1])
+    return subarrays
+
+
+def sort_k_increasing_decreasing_array(array: List[int]) -> List[int]:
+    subarrs = get_subarrays(array)
+    heap = [(arr[0], 0, i) for i, arr in enumerate(subarrs)]
+    heapq.heapify(heap)
+    merged = []
+    while heap:
+        val, idx, arr_id = heapq.heappop(heap)
+        merged.append(val)
+        new_idx = idx+1
+        parent_arr = subarrs[arr_id]
+        if new_idx < len(parent_arr):
+            heapq.heappush(heap, (parent_arr[new_idx], new_idx, arr_id))
+
+    return merged
 
 
 if __name__ == '__main__':
+    cases = [
+        # [1, 2, 3, 4, 5, 4, 3, 2, 1, 2, 3, 4, 5],
+        # [1, 2, 3, 2, 1, 4, 5, 10, 9, 4, 4, 1, -1],
+        # [1, 2, 4, 4, 5, 1, 2, 3, 4, 5, -10, 10, -10],
+        [2147483647, 8, 4, 2, 1, 0, -1, -2147483648]
+    ]
+    for case in cases:
+        # print(get_subarrays(case))
+        expected = sorted(case)
+        actual = sort_k_increasing_decreasing_array(case)
+        assert actual == expected, f"\ncase: {case}\n{actual} != {expected}"
     exit(
         generic_test.generic_test_main('sort_increasing_decreasing_array.py',
                                        'sort_increasing_decreasing_array.tsv',

elements-of-programming-interviews/python/sorted_arrays_merge.py

@@ -2,10 +2,22 @@
 
 from test_framework import generic_test
 
+import heapq
+
 
 def merge_sorted_arrays(sorted_arrays: List[List[int]]) -> List[int]:
-    # TODO - you fill in here.
-    return []
+    heap = [(arr[0], 0, i) for i, arr in enumerate(sorted_arrays) if arr]
+    heapq.heapify(heap)
+    merged = []
+    while heap:
+        val, idx, arr_id = heapq.heappop(heap)
+        merged.append(val)
+        new_idx = idx+1
+        parent_arr = sorted_arrays[arr_id]
+        if new_idx < len(parent_arr):
+            heapq.heappush(heap, (parent_arr[new_idx], new_idx, arr_id))
+
+    return merged
 
 
 if __name__ == '__main__':