Meta Coding Puzzles - Level 4 - Conveyor Chaos (Python)

conveyor_chaos.py

from typing import List
from sortedcontainers import SortedList
# Write any import statements here

def getMinExpectedHorizontalTravelDistance(N: int, H: List[int], A: List[int], B: List[int]) -> float:

    conveyors = get_conveyors(H, A, B)
    find_conveyors_below(conveyors)
    find_init_probs(conveyors)
    find_total_probs(conveyors)
    find_expected_distances(conveyors)
    find_delta_expected_distances(conveyors)
    total_expected_distance = get_total_expected_distance(conveyors)
    max_delta = max(conveyor.delta_expected_distance for conveyor in conveyors)

    return total_expected_distance - max_delta


class Conveyor:
    def __init__(self, height, left, right):
        self.height = height
        self.left = left
        self.right = right
        self.left_below = None
        self.right_below = None

        self.init_prob = 0.0
        self.prob = 0.0
        self.left_expected_distance = 0
        self.right_expected_distance = 0
        self.above_conveyors = []
        self.ceiling_intervals = []
        self.delta_expected_distance = 0

    def get_total_expected_distance(self):
        return 0.5 * (self.right - self.left) + 0.5 * self.left_expected_distance + 0.5 * self.right_expected_distance

    def __lt__(self, other):
        return self.height < other.height

    def __str__(self):
        return f'Conveyor: h: {self.height} l: {self.left} r: {self.right}, p: {self.prob}, delta: {self.delta_expected_distance}'


def get_conveyors(H: List[int], A: List[int], B: List[int]) -> List[Conveyor]:
    conveyors = []
    for i in range(len(H)):
        conveyors.append(Conveyor(H[i], A[i], B[i]))
    return conveyors


def find_conveyors_below(conveyors: List[Conveyor]) -> None:
    def left_key(endpoint):
        x, label, conveyor = endpoint
        if label == 'L':
            return x, 1
        else:
            return x, 0

    def right_key(endpoint):
        x, label, conveyor = endpoint
        if label == 'R':
            return -x, 1
        else:
            return -x, 0

    endpoints = sorted([(conveyor.left, 'L', conveyor) for conveyor in conveyors] +
                       [(conveyor.right, 'R', conveyor) for conveyor in conveyors],
                       key=left_key)
    tree = SortedList()

    buffer = []
    buffer_coord = None
    for endpoint in endpoints:
        x, label, conveyor = endpoint
        if label == 'R':
            for buff_conveyor in buffer:
                tree.add(buff_conveyor)
            buffer = []
            buffer_coord = None
            tree.remove(conveyor)
        else:
            if x != buffer_coord:
                for buff_conveyor in buffer:
                    tree.add(buff_conveyor)
                buffer = [conveyor]
                buffer_coord = x
            else:
                buffer.append(conveyor)
            conveyor_below = tree.bisect_left(conveyor) - 1
            if conveyor_below >= 0:
                conveyor.left_below = tree[conveyor_below]


    tree = SortedList()
    endpoints = sorted([(conveyor.left, 'L', conveyor) for conveyor in conveyors] +
                       [(conveyor.right, 'R', conveyor) for conveyor in conveyors],
                       key=right_key)

    buffer = []
    buffer_coord = None
    for endpoint in endpoints:
        x, label, conveyor = endpoint
        if label == 'L':
            for buff_conveyor in buffer:
                tree.add(buff_conveyor)
            buffer = []
            buffer_coord = None
            tree.remove(conveyor)
        else:
            if x != buffer_coord:
                for buff_conveyor in buffer:
                    tree.add(buff_conveyor)
                buffer = [conveyor]
                buffer_coord = x
            else:
                buffer.append(conveyor)
            conveyor_below = tree.bisect_left(conveyor) - 1
            if conveyor_below >= 0:
                conveyor.right_below = tree[conveyor_below]

def find_init_probs(conveyors: List[Conveyor], x_limit=1_000_000.0) -> None:
    def left_key(endpoint):
        x, label, conveyor = endpoint
        if label == 'L':
            return x, 1
        else:
            return x, 0

    conveyor_endpoints = sorted(([(conveyor.left, 'L', conveyor) for conveyor in conveyors] +
                                 [(conveyor.right, 'R', conveyor) for conveyor in conveyors]),
                                 key=left_key)
    tree = SortedList()
    curr_top_height = -1.0
    curr_top_beg = -1.0
    curr_top_conveyor = None
    for endpoint, label, conveyor in conveyor_endpoints:
        if label == 'L':
            if conveyor.height > curr_top_height:
                if curr_top_conveyor is not None:
                    curr_top_conveyor.init_prob += (endpoint - curr_top_beg) / x_limit
                    curr_top_conveyor.ceiling_intervals.append((curr_top_beg, endpoint))
                curr_top_height = conveyor.height
                curr_top_beg = endpoint
                curr_top_conveyor = conveyor
            tree.add(conveyor)

        else:
            tree.remove(conveyor)
            if conveyor.height == curr_top_height:
                conveyor.init_prob += (endpoint - curr_top_beg) / x_limit
                conveyor.ceiling_intervals.append((curr_top_beg, endpoint))
                curr_top_conveyor = tree[-1] if len(tree) > 0 else None
                curr_top_height = curr_top_conveyor.height if curr_top_conveyor is not None else -1
                curr_top_beg = endpoint


def find_total_probs(conveyors: List[Conveyor]) -> None:
    conveyors.sort(key=lambda conveyor: -conveyor.height)
    for conveyor in conveyors:
        conveyor.prob = conveyor.init_prob

    for conveyor in conveyors:
        if conveyor.left_below is not None:
            conveyor.left_below.prob += 0.5 * conveyor.prob
        if conveyor.right_below is not None:
            conveyor.right_below.prob += 0.5 * conveyor.prob


def find_expected_distances(conveyors: List[Conveyor]) -> None:
    conveyors.sort(key=lambda conveyor: conveyor.height)

    for conveyor in conveyors:
        if conveyor.left_below is not None:
            conveyor.left_expected_distance += conveyor.left_below.get_total_expected_distance()
        if conveyor.right_below is not None:
            conveyor.right_expected_distance += conveyor.right_below.get_total_expected_distance()


def find_delta_expected_distances(conveyors: List[Conveyor], x_limit=1_000_000.0) -> None:
    for conveyor in conveyors:
        if conveyor.left_below is not None:
            conveyor.left_below.above_conveyors.append((0.5 * conveyor.prob, conveyor.left))
        if conveyor.right_below is not None:
            conveyor.right_below.above_conveyors.append((0.5 * conveyor.prob, conveyor.right))

    for conveyor in conveyors:
        expected_distance_left, expected_distance_right = 0.0, 0.0

        for beg, end in conveyor.ceiling_intervals:
            prob = (end - beg) / x_limit
            expected_distance_left += prob * (beg - conveyor.left + (end-beg)/2 + conveyor.left_expected_distance)
            expected_distance_right += prob * (conveyor.right - end + (end-beg)/2 + conveyor.right_expected_distance)

        for prob, x in conveyor.above_conveyors:
            expected_distance_left += prob * (x - conveyor.left + conveyor.left_expected_distance)
            expected_distance_right += prob * (conveyor.right - x + conveyor.right_expected_distance)

        min_expected_distance = min(expected_distance_left, expected_distance_right)
        expected_distance = 0.5 * (expected_distance_left + expected_distance_right)
        conveyor.delta_expected_distance = abs(min_expected_distance - expected_distance)


def get_total_expected_distance(conveyors: List[Conveyor]) -> float:
    total_expected_distance = 0.0
    for conveyor in conveyors:
        if len(conveyor.ceiling_intervals) > 0:
            total_expected_distance += conveyor.get_total_expected_distance() * conveyor.init_prob
    return total_expected_distance