Challenge - Taco Truck

You own a taco truck that’s open 24/7 and manage five employees who run it. Employees work solo, eight-hour shifts. You decide the best way to set their schedule for the upcoming week is to create a bunch of random schedules and select one that looks best.

You build a 1000x21 array of random employee ids where element (i,j) gives the employee id working shift j for schedule i.

A Schedule is valid as long as no employee works two consecutive shifts. Get the row indices of all valid schedules.

import numpy as np

generator = np.random.default_rng(999)
schedules = generator.integers(low=0, high=5, size=(1000, 21))

print(schedules)
# [[4 3 0 ... 2 0 0]
#  [2 4 3 ... 3 3 2]
#  [1 0 1 ... 1 2 1]
#  ...
#  [2 2 1 ... 3 1 4]
#  [1 0 3 ... 2 3 2]
#  [1 1 4 ... 2 4 2]]

[[4 3 0 ... 2 0 0]
 [2 4 3 ... 3 3 2]
 [1 0 1 ... 1 2 1]
 ...
 [2 2 1 ... 3 1 4]
 [1 0 3 ... 2 3 2]
 [1 1 4 ... 2 4 2]]

employees = np.unique(schedules)
argsorted_schedules = np.argsort(schedules, axis=1) 
sorted_schedules = np.sort(schedules, axis=1)
is_diff = np.full((1000), False)

for employee in employees:
    selected_employee = sorted_schedules == employee
    argsorted_schedules_copy = argsorted_schedules.copy()
    argsorted_schedules_copy.ravel()[(~selected_employee).ravel()] = -100
    argsorted_schedules_copy = np.sort(argsorted_schedules_copy)
    differences = np.any(np.diff(argsorted_schedules_copy, axis=1) == 1, axis=1)
    is_diff = np.logical_or(is_diff, differences)

np.where(~is_diff)

(array([ 25, 138, 188, 289, 375, 426, 533, 886, 975, 982]),)

is_valid = np.all(schedules[:, :-1] != schedules[:, 1:], axis=1)
np.nonzero(is_valid)

(array([ 25, 138, 188, 289, 375, 426, 533, 886, 975, 982]),)