Generalized find K elements in an array that adds up to a sum

find_k_elements_that_sums_to.py

# track number of iterations
BIG_O = 0

def find_k_numbers_that_add_upto_sum(arr, k, sum):
    """
    A generic solution to the problem of Finding k numbers in an array
    that add up to a given sum.
    The solution is based on the assumption that the numbers in the array
    are distinct.
    - Works for both positive and negative numbers.
    - Works for any k.
    - Worse time complexity of O(N * K) where N is the length of the array
    K is the number of elements we are looking for.
    - If there are multiple solutions, return first one of them.
    - If no such combination exists, return an empty list.
    """
    array_len = len(arr)
    if array_len < k or k < 1:
        return []
    
    # store available numbers in a map so we can check in O(1) if it is
    # the remainder we need to add to make up to target sum
    numbers_map = {}
    for num in arr:
        if numbers_map.get(num) != None:
            numbers_map[num] += 1
        else:
            numbers_map[num] = 1
    
    res = find_k_numbers_that_add_upto_sum_helper(arr, k, sum, numbers_map)
    return res

def find_k_numbers_that_add_upto_sum_helper(arr, k, sum, numbers_map):
    global BIG_O
    BIG_O += 1

    # print(arr, "<---")
    if len(arr) < 1:
        return []
    
    first_number = arr[0]
    remainder = sum - first_number

    if k == 1:
        if first_number == sum:
            return [first_number]
        return []

    if k == 2:
        if check_if_number_is_addable(remainder, numbers_map):
            return [first_number, remainder]


    check_if_number_is_addable(first_number, numbers_map)
    # check from arr[1:] to see if we can find candidates to add to
    # arr[0] to give us sum
    res = find_k_numbers_that_add_upto_sum_helper(
        shift_array_forward(arr), k - 1, remainder, numbers_map
    )

    # found sub-array of length k - 1 that its sum + first_number == target sum
    if res and len(res) == k - 1:
        return [first_number, *res]
    
    return find_k_numbers_that_add_upto_sum_helper(
        shift_array_forward_by_k(arr, k - 1), k, sum, numbers_map
    )

def shift_array_forward_by_k(arr, k):
    """ resize array in O(k) by removing first k items """
    for i in range(k):
        remmove_element_at_index(arr, i)
    return arr

def remmove_element_at_index(arr, index):
    """ resize array in O(k) by removing first k items """
    if len(arr) < index + 1:
        return arr
    arr[index] = arr[len(arr) - 1]
    arr.pop()
    return arr

def shift_array_forward(arr):
    """ resize array in O(1) by removing first item """
    arr[0] = arr[len(arr) - 1]
    arr.pop()
    return arr

def check_if_number_is_addable(num, numbers_map):
    """ check if num is the remainder we need to add to make up to target sum """
    if numbers_map.get(num) != None:
        if numbers_map[num] > 0:
            numbers_map[num] -= 1
            return True
    return False


if __name__ == "__main__":
    arr = [2, -2, 1, 6, 10, 26, -97, 100]

    sum = 1
    k = 3

    res = find_k_numbers_that_add_upto_sum(arr, k, sum)

    print()
    if len(res):
        print("sum of:", res)
        print("is:", sum)
        print()
    else:
        print("not found")
    print()
    print(BIG_O, "number of iterations -- BIG O time complexity")