Meta Coding Puzzles - Level 4 - Mathematical Art (C++)

mathematical_art.cpp

#include <bits/stdc++.h>
using namespace std;
// Write any include statements here

#define MAXN 4000200

long long bit[2 * MAXN];

void bit_upd(int idx, long long v) {
    while (idx < 2 * MAXN) {
        bit[idx] += v;
        idx += idx & -idx;
    }
}

long long bit_query(int idx) {
    long long res = 0;
    while (idx > 0) {
        res += bit[idx];
        idx -= idx & -idx;
    }
    return res;
}

struct Line {
    long long coord, beg, end;

    Line(long long c, long long b, long long e) : coord(c), beg(b), end(e) {}
};

tuple<vector<Line>, vector<Line>> get_lines(vector<int> L, string D) {
    vector<Line> horizontal_lines, vertical_lines;
    long long x = 0, y = 0;
    for (int i = 0; i < L.size(); i++) {
        long long d = D[i];
        long long l = L[i];

        if (d == 'U') {
            vertical_lines.push_back(Line(x, y, y+l));
            y += l;
        }
        else if (d == 'D') {
            vertical_lines.push_back(Line(x, y-l, y));
            y -= l;
        }
        else if (d == 'R') {
            horizontal_lines.push_back(Line(y, x, x+l));
            x += l;
        }
        else if (d == 'L') {
            horizontal_lines.push_back(Line(y, x-l, x));
            x -= l;
        }
    }
    return {horizontal_lines, vertical_lines};
}

vector<Line> get_merged_lines(vector<Line> lines) {
    vector<Line> merged_lines;
    map<long long, vector<pair<long long , long long>>> coord_lines;

    for (auto line : lines) {
        coord_lines[line.coord].push_back({line.beg, line.end});
    }

    for (auto [coord, lines] : coord_lines) {
        sort(lines.begin(), lines.end());
        long long curr_beg = lines[0].first;
        long long curr_end = lines[0].second;
        for (int i=1; i<lines.size(); i++) {
            long long next_beg = lines[i].first;
            long long next_end = lines[i].second;

            if (next_beg <= curr_end) {
                curr_end = max(curr_end, next_end);
            }
            else {
                merged_lines.push_back(Line(coord, curr_beg, curr_end));
                curr_beg = next_beg;
                curr_end = next_end;
            }
        }
        merged_lines.push_back(Line(coord, curr_beg, curr_end));
    }

    return merged_lines;
}

map<long long, long long> get_compressed_map(vector<long long> coords) {
    sort(coords.begin(), coords.end());
    auto it = unique(coords.begin(), coords.end());
    coords.resize(distance(coords.begin(), it));

    map<long long, long long> compressed_map;
    for (int i=0; i<coords.size(); i++) {
        compressed_map[coords[i]] = i+1;
    }
    return compressed_map;
}

tuple<vector<Line>, vector<Line>> get_compressed_lines(vector<Line> hor_lines, vector<Line> ver_lines) {
    vector<long long> x_coords, y_coords;

    for (auto line : hor_lines) {
        x_coords.push_back(line.beg);
        x_coords.push_back(line.end);
        y_coords.push_back(line.coord);
    }

    for (auto line : ver_lines) {
        y_coords.push_back(line.beg);
        y_coords.push_back(line.end);
        x_coords.push_back(line.coord);
    }

    auto x_map = get_compressed_map(x_coords);
    auto y_map = get_compressed_map(y_coords);

    vector<Line> compr_hor_lines, compr_ver_lines;

    for (auto line : hor_lines) {
        compr_hor_lines.push_back(Line(y_map[line.coord], x_map[line.beg], x_map[line.end]));
    }

    for (auto line : ver_lines) {
        compr_ver_lines.push_back(Line(x_map[line.coord], y_map[line.beg], y_map[line.end]));
    }

    return {compr_hor_lines, compr_ver_lines};
}

bool compare_coord(Line a, Line b) {
    return a.coord < b.coord;
}

bool compare_beg(Line a, Line b) {
    return a.beg < b.beg;
}

long long getPlusSignCount(int N, vector<int> L, string D) {

  auto [horizontal_lines, vertical_lines] = get_lines(L, D);

    auto merged_hor_lines = get_merged_lines(horizontal_lines);
    auto merged_ver_lines = get_merged_lines(vertical_lines);

    auto [compr_hor_lines, compr_ver_lines] = get_compressed_lines(merged_hor_lines, merged_ver_lines);

    sort(compr_hor_lines.begin(), compr_hor_lines.end(), compare_coord);
    sort(compr_ver_lines.begin(), compr_ver_lines.end(), compare_beg);

    long long n_plus_signs = 0;
    priority_queue<pair<long long, long long>,
                   vector<pair<long long, long long>>,
                   greater<pair<long long, long long>>> pq;
    int ptr_ver_lines = 0;
    for (auto hor_line : compr_hor_lines) {

        while (ptr_ver_lines < compr_ver_lines.size()) {
            auto ver_line = compr_ver_lines[ptr_ver_lines];
            if (ver_line.beg < hor_line.coord) {
                ptr_ver_lines++;
                pq.push({ver_line.end, ver_line.coord});
                bit_upd(ver_line.coord, 1);
            }
            else break;
        }

        while (pq.size() > 0) {
            auto [ver_line_end, ver_line_coord] = pq.top();
            if (ver_line_end <= hor_line.coord) {
                pq.pop();
                bit_upd(ver_line_coord, -1);
            }
            else break;
        }

        if (hor_line.beg + 1 <= hor_line.end - 1) {
            n_plus_signs += bit_query(hor_line.end - 1) - bit_query(hor_line.beg);
        }
    }

    return n_plus_signs;
}