Algorithm to position diners at a table with social distancing guidelines

Trying to solve this task:

A cafeteria table consists of a row of N seats, numbered from 1 to N from left to right. Social distancing guidelines require that every diner be seated such that K seats to their left and K seats to their right (or all the remaining seats to that side if there are fewer than K) remain empty. There are currently M diners seated at the table, the ith of whom is in seat S(i).

No two diners are sitting in the same seat, and the social distancing guidelines are satisfied. Determine the maximum number of additional diners who can potentially sit at the table without social distancing guidelines being violated for any new or existing diners, assuming that the existing diners cannot move and that the additional diners will cooperate to maximize how many of them can sit down. Please take care to write a solution which runs within the time limit.

Constraints:

1 <= N <= 10^{15} 
1 <= K <=  N
1 <= M <= 500000
M <= N
1 <= S(i) <= N

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

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

Sample Explanation

In the first case, the cafeteria table has N=10 seats, with two diners currently at seats 2 and 6 respectively. The table initially looks as follows, with brackets covering the K=1 seat to the left and right of each existing diner that may not be taken.

  1 2 3 4 5 6 7 8 9 10
  [   ]   [   ]

Three additional diners may sit at seats 4, 8, and 10 without violating the social distancing guidelines. In the second case, only 1 additional diner is able to join the table, by sitting in any of the first 3 seats.

My solution works for both test cases (1 and 2):

def getMaxAdditionalDinersCount(N: int, K: int, M: int, S: List[int]) -> int:
        
    if N == 0:
        return 0
    
    if K == 0:
        return N
    
    if not S or M == 0:
        return N // (K+1)
    
    pos_cnts = 0   # updated position counts    
    l = sorted(S)  # sort positions in increasing order
   
    first_used_pos = l[0]
    last_used_pos = l[len(l)-1]
    max_pos = N
    tail_len = max_pos - last_used_pos

    # if zero position is not taken yet
    if (1 not in S):
        new_pos_cnt  = first_used_pos // (K+1) - 1
        pos_cnts += new_pos_cnt # update count of all positions
     
    # if last position is not taken yet
    if max_pos not in S:
        new_pos_cnt  = tail_len // (K+1)
        pos_cnts += new_pos_cnt # update count of all positions
    
    indx_prev = l[0]  # index of previous position 
    for indx in l[1:]: # iterate existing positions
        gap = indx - indx_prev
        new_pos_cnt = gap // (K+1) - 1
        pos_cnts += new_pos_cnt # update count of all positions
        indx_prev = indx        
            
    return pos_cnts

Yet for a test case:

N = 10
K = 1
M = 2
S = [3,5]

It returns the wrong answer 2 instead of 3, not taking into account free position 1. What would be the correct algorithm?

Answer

The issue is within the if condition that checks the first position:

# if zero position is not taken yet
if (1 not in S):
    new_pos_cnt  = first_used_pos // (K+1) - 1
    pos_cnts += new_pos_cnt # update count of all positions

The problem is that you have to treat it differently when there is a remainder in the division. You need to modify the new_pos_cnt calculation so that it doesn’t subtract a position if there is a remainder:

# if zero position is not taken yet
if (1 not in S):
    if first_used_pos // (K+1) == first_used_pos / (K+1):
       new_pos_cnt  = first_used_pos // (K+1) - 1
    else:
       new_pos_cnt  = first_used_pos // (K+1) 
    pos_cnts += new_pos_cnt # update count of all positions  

This produces the same results for your first two test cases and the correct result for the third case as well.

For the testcase 1:
first_used_pos = 2
K = 1
first_used_pos // (K+1) -->  2 // (1+1) = 1
first_used_pos /  (K+1) -->  2 /  (1+1) = 1

For the testcase 2:
first_used_pos = 6
K = 2
first_used_pos // (K+1) -->  6 // (2+1) = 2
first_used_pos /  (K+1) -->  6 /  (2+1) = 2

For the testcase 3:
first_used_pos = 3
K = 1
first_used_pos // (K+1) -->  3 // (1+1) = 1
first_used_pos /  (K+1) -->  3 /  (1+1) = 1.5 # Remainder exist -> don't subtract