Meta Coding Puzzles - Level 4 - Conveyor Chaos (C++)

conveyor_chaos.cpp

#include <iostream>
#include <string>
#include <set>
#include <vector>
using namespace std;
// Write any include statements here

struct Conveyor {
    int height, left, right;
    double init_prob, prob;
    double left_expected_distance, right_expected_distance, delta_expected_distance;

    vector<pair<double, int>> above_conveyors;
    vector<pair<int,int>> ceiling_intervals;

    Conveyor* left_below;
    Conveyor* right_below;

    Conveyor(int h, int l, int r) : height(h), left(l), right(r) {
        init_prob = 0.0;
        prob = 0.0;
        left_expected_distance = 0.0;
        right_expected_distance = 0.0;
        delta_expected_distance = 0.0;
        left_below = nullptr;
        right_below = nullptr;
    }

    double get_total_expected_distance() {
        return 0.5 * (right - left + left_expected_distance + right_expected_distance);
    }

    bool operator<(const Conveyor& other) const {
        return height < other.height;
    }

    bool operator>(const Conveyor& other) const {
        return height > other.height;
    }

    friend ostream& operator<<(ostream& os, const Conveyor& conveyor) {
        os << "Conveyor(" << conveyor.height << ", " << conveyor.left << ", " << conveyor.right << ", " <<
            conveyor.init_prob << ", " << conveyor.prob << ", " <<
                conveyor.left_expected_distance << ", " << conveyor.right_expected_distance << ", " <<
                    conveyor.delta_expected_distance << ")";
        return os;
    }
};

struct Endpoint {
    int x;
    string label;
    Conveyor* conveyor;

    friend ostream& operator<<(ostream& os, const Endpoint& endpoint) {
        os << "Endpoint(" << endpoint.x << ", " << endpoint.label << ", " << endpoint.conveyor << ")";
        return os;
    }
};

// Custom comparator for sorting by height in descending order
struct CompareConveyorPtrs {
    bool operator()(const Conveyor* a, const Conveyor* b) const {
        return a->height < b->height; // Sort in descending order by height
    }
};

bool cmp_ptrs_lower(Conveyor* a, Conveyor* b) {
    return a->height < b->height;
}

bool cmp_ptrs_greater(Conveyor* a, Conveyor* b) {
    return a->height > b->height;
}

vector<Conveyor> get_conveyors(vector<int> H, vector<int> A, vector<int> B) {
    vector<Conveyor> conveyors;
    for (int i = 0; i < H.size(); i++) {
        conveyors.push_back(Conveyor(H[i], A[i], B[i]));
    }
    return conveyors;
}

bool cmp_left(Endpoint &a, Endpoint &b) {
    if (a.x == b.x) return (a.label == "right" and b.label == "left");
    else return a.x < b.x;
}

bool cmp_right(Endpoint &a, Endpoint &b) {
    if (a.x == b.x) return (a.label == "left" and b.label == "right");
    else return a.x > b.x;
}

bool find_conveyors_below(vector<Conveyor>& conveyors) {
    vector<Endpoint> endpoints;
    for (auto &conveyor : conveyors) {
        endpoints.push_back({conveyor.left, "left", &conveyor});
        endpoints.push_back({conveyor.right, "right", &conveyor});
    }
    sort(endpoints.begin(), endpoints.end(), cmp_left);
    set<Conveyor*, CompareConveyorPtrs> tree;
    vector<Conveyor*> buffer;
    int buffer_coord = -1;

    for (auto &endpoint : endpoints) {
        int x = endpoint.x;
        string label = endpoint.label;
        Conveyor* conveyor = endpoint.conveyor;

        if (label == "right") {
            while (!buffer.empty()) {
                tree.insert(buffer.back());
                buffer.pop_back();
            }
            buffer_coord = -1;
            tree.erase(conveyor);
        } else {
            if (x != buffer_coord) {
                while (!buffer.empty()) {
                    tree.insert(buffer.back());
                    buffer.pop_back();
                }
                buffer_coord = x;
            }
            buffer.push_back(conveyor);

            auto conveyor_below_it = tree.lower_bound(conveyor);
            if (conveyor_below_it != tree.begin()) {
                --conveyor_below_it;
                conveyor->left_below = *conveyor_below_it;
            }
        }
    }

    sort(endpoints.begin(), endpoints.end(), cmp_right);
    tree.clear();
    buffer.clear();
    buffer_coord = -1;

    for (auto &endpoint : endpoints) {
        int x = endpoint.x;
        string label = endpoint.label;
        Conveyor* conveyor = endpoint.conveyor;

        if (label == "left") {
            while (!buffer.empty()) {
                tree.insert(buffer.back());
                buffer.pop_back();
            }
            tree.erase(conveyor);
            buffer_coord = -1;
        } else {
            if (x != buffer_coord) {
                while (!buffer.empty()) {
                    tree.insert(buffer.back());
                    buffer.pop_back();
                }
                buffer_coord = x;
            }
            buffer.push_back(conveyor);

            auto conveyor_below_it = tree.lower_bound(conveyor);
            if (conveyor_below_it != tree.begin()) {
                --conveyor_below_it;
                conveyor->right_below = *conveyor_below_it;
            }
        }
    }
  return false;
}

void find_init_probs(vector<Conveyor>& conveyors) {
    double x_lim = 1e6;
    vector<Endpoint> endpoints;
    for (auto &conveyor : conveyors) {
        endpoints.push_back({conveyor.left, "left", &conveyor});
        endpoints.push_back({conveyor.right, "right", &conveyor});
    }
    sort(endpoints.begin(), endpoints.end(), cmp_left);

    set<Conveyor*, CompareConveyorPtrs> tree;
    int curr_top_height = -1;
    int curr_top_beg = -1;
    Conveyor* curr_top_conveyor = nullptr;

    for (auto &endpoint : endpoints) {
        int x = endpoint.x;
        string label = endpoint.label;
        Conveyor* conveyor = endpoint.conveyor;

        if (label == "left") {
            if (conveyor->height > curr_top_height) {
                if (curr_top_conveyor != nullptr) {
                    curr_top_conveyor->ceiling_intervals.push_back({curr_top_beg, x});
                    curr_top_conveyor->init_prob += double(x - curr_top_beg) / x_lim;
                }
                curr_top_height = conveyor->height;
                curr_top_beg = x;
                curr_top_conveyor = conveyor;
            }
            tree.insert(conveyor);
        } else {
            tree.erase(conveyor);
            if (conveyor->height == curr_top_height) {
                conveyor->init_prob += double(x - curr_top_beg) / x_lim;
                conveyor->ceiling_intervals.push_back({curr_top_beg, x});
                if (tree.empty()) {
                    curr_top_height = -1;
                    curr_top_beg = -1;
                    curr_top_conveyor = nullptr;
                } else {
                    curr_top_conveyor = *tree.rbegin();
                    curr_top_height = curr_top_conveyor->height;
                    curr_top_beg = x;
                }
            }
        }
    }
}

void find_total_probs(vector<Conveyor>& conveyors) {
    vector<Conveyor*> conveyor_ptrs;
    for (auto &conveyor : conveyors) {
        conveyor_ptrs.push_back(&conveyor);
    }
    sort(conveyor_ptrs.begin(), conveyor_ptrs.end(), cmp_ptrs_greater);

    for (auto &conveyor : conveyors) {
        conveyor.prob = conveyor.init_prob;
    }

    for (auto conveyor : conveyor_ptrs) {
        if (conveyor->left_below != nullptr) {
            conveyor->left_below->prob += 0.5 * conveyor->prob;
        }
        if (conveyor->right_below != nullptr) {
            conveyor->right_below->prob += 0.5 * conveyor->prob;
        }
    }
}

void find_expected_distances(vector<Conveyor>& conveyors) {
    vector<Conveyor*> conveyor_ptrs;
    for (auto &conveyor : conveyors) {
        conveyor_ptrs.push_back(&conveyor);
    }
    sort(conveyor_ptrs.begin(), conveyor_ptrs.end(), cmp_ptrs_lower);

    for (auto conveyor : conveyor_ptrs) {
        if (conveyor->left_below != nullptr) {
            conveyor->left_expected_distance += conveyor->left_below->get_total_expected_distance();
        }
        if (conveyor->right_below != nullptr) {
            conveyor->right_expected_distance += conveyor->right_below->get_total_expected_distance();
        }
    }
}

void find_delta_expected_distances(vector<Conveyor> &conveyors) {
double x_lim = 1e6;

    for (auto &conveyor : conveyors) {
        if (conveyor.left_below != nullptr) {
            conveyor.left_below->above_conveyors.push_back({0.5 * conveyor.prob, conveyor.left});
        }
        if (conveyor.right_below != nullptr) {
            conveyor.right_below->above_conveyors.push_back({0.5 * conveyor.prob, conveyor.right});
        }
    }

    for (auto &conveyor : conveyors) {
        double expected_distance_left=0.0, expected_distance_right=0.0;

        for (auto [beg, end] : conveyor.ceiling_intervals) {
            double interval_length = end - beg;
            double prob = interval_length / x_lim;
            expected_distance_left += prob * (beg - conveyor.left + 0.5 * interval_length + conveyor.left_expected_distance);
            expected_distance_right += prob * (conveyor.right - end + 0.5 * interval_length + conveyor.right_expected_distance);
        }

        for (auto [prob, x] : 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);
        }

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

double get_total_expected_distance(vector<Conveyor>& conveyors) {
    double total_expected_distance = 0.0;
    for (auto conveyor : conveyors) {
        if (!conveyor.ceiling_intervals.empty()) {
            total_expected_distance += conveyor.get_total_expected_distance() * conveyor.init_prob;
        }
    }
    return total_expected_distance;
}

double getMinExpectedHorizontalTravelDistance(int N, vector<int> H, vector<int> A, vector<int> B) {
    vector<Conveyor> conveyors = get_conveyors(H, A, B);

    if (find_conveyors_below(conveyors)) return -1;
  
    find_init_probs(conveyors);

    find_total_probs(conveyors);
    find_expected_distances(conveyors);
    find_delta_expected_distances(conveyors);

    double total_expected_distance = get_total_expected_distance(conveyors);
    double max_delta = 0;
    for (auto conveyor : conveyors) {
        max_delta = max(max_delta, conveyor.delta_expected_distance);
    }

    return total_expected_distance - max_delta;
}