elevator/elevator.py

#!/usr/bin/env python3

from enum import Enum

class Elevator:
	def __init__(self, e_id, floor, goal):
		self.id = e_id
		self.floor = floor
		self.goal = goal
		self.dropoffs = []

	def direction(self):
		if self.goal is None:
			return None
		if self.goal > self.floor:
			return Direction.UP
		else:
			return Direction.DOWN

	def dropoff(self, floor):
		if floor in self.dropoffs:
			assert self.goal is not None
			return
		assert floor != self.floor

		if self.goal is None:
			self.goal = floor
			self.dropoffs.append(floor)
		else:
			direction = self.direction()
			if direction == Direction.UP:
				if floor > self.goal:
					self.goal = floor
				self.dropoffs.append(floor)
				self.dropoffs.sort()
			elif direction == Direction.DOWN:
				if floor < self.goal:
					self.goal = floor
				self.dropoffs.append(floor)
				self.dropoffs.sort(reverse=True)

	def __repr__(self):
		attrs = ', '.join('%s=%s' % t for t in self.__dict__.items())
		return 'Elevator(%s)' % attrs

class Direction(Enum):
	UP = 1
	DOWN = 2

class ElevatorControlSystem:
	def __init__(self, num_elevators):
		self.elevators = {i: Elevator(i, 1, None) for i in range(num_elevators)}
		self.pickup_requests = []

	def pickup(self, floor, direction):
		self.pickup_requests.append((floor, direction))

	def dropoff(self, elevator_id, floor):
		self.elevators[elevator_id].dropoff(floor)

	def step(self):
		# serve pickup requests
		free_elevators = list(filter(lambda e: e.goal is None, self.elevators.values()))
		while free_elevators and self.pickup_requests:
			floor, direction = self.pickup_requests.pop(0)
			elevator = min(free_elevators, key=lambda e: abs(e.floor - floor))
			elevator.goal = floor

		# move elevators to requested floors
		for elevator in self.elevators.values():
			if elevator.goal is not None:
				if elevator.goal > elevator.floor:
					elevator.floor += 1
				elif elevator.goal < elevator.floor:
					elevator.floor -= 1
				# check dropoffs and goal after moving
				if elevator.dropoffs and elevator.dropoffs[0] == elevator.floor:
					elevator.dropoffs = elevator.dropoffs[1:]
				if elevator.goal == elevator.floor:
					elevator.goal = None

if __name__ == '__main__':
	ecs = ElevatorControlSystem(2)
	print(' '*4, ecs.elevators)

	print('requesting pickup on floor 2 down')
	ecs.pickup(2, Direction.DOWN)
	ecs.step()
	print(' '*4, ecs.elevators)

	print('requesting dropoff on floor 1 and pickup on floor 1 up')
	ecs.dropoff(0, 1)
	ecs.pickup(1, Direction.UP)
	ecs.step()
	print(' '*4, ecs.elevators)

	print('requesting dropoff on floor 3')
	ecs.dropoff(1, 3)
	ecs.step()
	print(' '*4, ecs.elevators)

	print('requesting pickup on floor 2 up')
	ecs.pickup(2, Direction.UP)
	ecs.step()
	print(' '*4, ecs.elevators)