n-hotels

In [2]:

    

from typing import List


    

In [46]:

    

def n_hotels(a: List[float]) -> List[int]:
    # b holds the best penalty for hotel i at b[i]
    b = [0] * len(a)
    b[0] = (200 - a[0]) ** 2
    
    # p holds the best previous hotel
    p = [0] * len(a)
    # -1 means there's no hotel before, we've jumped here from the start position
    p[0] = -1
    
    for i in range(1, len(a)):
        b[i], best_j = min([
            (b[j] + (200 - (a[i] - a[j])) ** 2, j)
            for j in range(i)
        ])
        p[i] = best_j
        # can just jump here from start directly
        if (200 - a[i]) ** 2 < b[i]:
            b[i] = (200 - a[i]) ** 2
            p[i] = -1
        
    tmp = p[-1]
    path = [tmp]
    while True:
        tmp = p[tmp]
        if tmp != -1:
            path.append(tmp)
        else:
            break
        
    return path[::-1] + [len(a) - 1]


    

In [47]:

    

assert n_hotels([200, 400, 600]) == [0, 1, 2]
assert n_hotels([200, 300, 400]) == [0, 2]
assert n_hotels([100, 200, 400]) == [1, 2]