Python: return value in a list if there are >= 2 values to left or right

Question

I am trying to calculate how many more people can sit at a table with N seats in a row. Social distancing requires K spaces between seats. M represents the number of people currently seated at the table, and S represents their seat number.

Example 1:

N = 10
K = 1
M = 2
S = [2, 6]
Expected Return Value = 3

Reasoning:

(Open seats represented by corresponding numbers, K represented as -, S represented as *)

Row: - * - 4 - * - 8 9 10     

Seats 4, 8, and 10 are open with 1 space before and/or after the seat, or the row begins/ends.

Example 2:

N = 15
K = 2
M = 3
S = [11, 6, 14]
Expected Return Value = 1

Reasoning:

(Open seats represented by corresponding numbers, K represented as -, S represented as *)

Row: 1 2 3 - - * - - - - * - - * -

Seat 3 is available, with 2 spaces before and/or after other seats, or the row begins/ends.

I am not sure how to account for K spaces before/after occupied seats ( S ), and/or the row ending.

Answer 1

Here an example for second input. It's a simple greedy solution. At first we mark all unavailable seats. Then we check every seat starting from the first if we can seat here

N = 15
K = 2
M = 3
S = [11, 6, 14]

unavailable_seats = set()


def reserve_seat(reserved_seat: int):
    # keeping soc distancing
    for seat in range(reserved_seat - K, reserved_seat + K + 1):
        # be sure we are not out of range
        if 1 <= seat <= N:
            # reserve seat
            unavailable_seats.add(seat)


# reserve all reserved seats
for reserved_seat in S:
    reserve_seat(reserved_seat)

new_seats_count = 0
for new_seat in range(1, N + 1):
    # try to reserve every seat if it is available
    if new_seat not in unavailable_seats:
        new_seats_count += 1
        reserve_seat(new_seat)

print(new_seats_count)