inserts the pins to columns with least consumed heights

"""
Given a set of pins with attributes (pin id, height), write a function that takes the argument 'k' (determining the number of columns)
and inserts the pins such that every pin goes into a column with least consumed height. If there is a tie then insert into the left most column.
"""

def model_pins(pins, n_columns):
    agg_pins = {}
    grid_layout = []
    for i in range(n_columns):
        grid_layout.append([])
        agg_pins[i] = 0

    for i in range(n_columns):  # fill first row for all columns
        agg_pins[i] += pins[i]['height']
        grid_layout[i] = [pins[i]]

    for i in range(n_columns, len(pins)):
        lowest_height_val = float('Inf')
        aux_index = -1
        for index in range(len(agg_pins)):  # find the column index with lowest height
            if agg_pins[index] < lowest_height_val:  # this covers the most left index in case of a tie
                lowest_height_val = agg_pins[index]
                aux_index = index

        agg_pins[aux_index] += pins[i]['height']
        grid_layout[aux_index].append(pins[i])

    return grid_layout



def main():
    pins = [
        {"id": 1, "height": 300},
        {"id": 2, "height": 200},
        {"id": 3, "height": 250},
        {"id": 4, "height": 350},
        {"id": 5, "height": 100},
    ]
    n_columns = 2

    expected_result = [
        [
            {"id": 1, "height": 300},
            {"id": 4, "height": 350},
        ],
        [
            {"id": 2, "height": 200},
            {"id": 3, "height": 250},
            {"id": 5, "height": 100},
        ]
    ]

    result = model_pins(pins, n_columns)
    assert result == expected_result