Meta Coding Puzzles - Level 4 - Mathematical Art (Python)

mathematical_art.py

from typing import List
from heapq import heappush, heappop
# Write any import statements here

def getPlusSignCount(N: int, L: List[int], D: str) -> int:

    # Write your code here
    horizontal_lines, vertical_lines = get_lines(L, D)

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

    comp_hor_lines, comp_ver_lines, max_x, max_y = get_compressed_lines(merged_hor_lines, merged_ver_lines)
    
    segtree = FenwickTree(0, max_x)

    comp_hor_lines = sorted(comp_hor_lines, key=lambda line: line[0])
    comp_ver_lines = sorted(comp_ver_lines, key=lambda line: line[1][0])

    n_plus_signs = 0
    priority_queue = []

    for y, (x_beg, x_end) in comp_hor_lines:

        while len(comp_ver_lines) > 0:
            x, (y_beg, y_end) = comp_ver_lines[0]
            if y_beg < y:
                comp_ver_lines.pop(0)
                heappush(priority_queue, (y_end, x))

                segtree.update(x, 1)
            else:
                break

        while len(priority_queue) > 0:
            y_end, x = priority_queue[0]

            if y_end <= y:
                heappop(priority_queue)
                segtree.update(x, -1)
            else:
                break

        if x_beg+1 <= x_end-1:
            n_plus_signs += segtree.query(x_beg+1, x_end-1)
            
    return -1
    
    return n_plus_signs


class FenwickTree:
    def __init__(self, left, right):
        self.left = left
        self.right = right
        self.nodes = [0 for _ in range(right+10)]

    def update(self, idx, val):
        while idx <= self.right:
            self.nodes[idx] += val
            idx += idx & -idx

    def query(self, beg, end):
        return self._query(end) - self._query(beg-1)

    def _query(self, idx):
        total = 0
        while idx > 0:
            total += self.nodes[idx]
            idx -= idx & -idx
        return total

def get_lines(L: List[int], D: str):
    horizontal_lines = []
    vertical_lines = []

    x, y = 0, 0
    for l, d in zip(L, D):
        if d == 'U':
            vertical_lines.append((x, (y, y + l)))
            y = y + l
        if d == 'D':
            vertical_lines.append((x, (y - l, y)))
            y = y - l
        if d == 'R':
            horizontal_lines.append((y, (x, x + l)))
            x = x + l
        if d == 'L':
            horizontal_lines.append((y, (x - l, x)))
            x = x - l

    return horizontal_lines, vertical_lines


def get_compressed_coords(coords):
    coords = sorted(set(coords))
    coords_map = {coord: idx+1 for idx, coord in enumerate(coords)}
    return coords_map


def get_compressed_lines(horizontal_lines, vertical_lines):
    x_coords, y_coords = [], []
    for y, (x1, x2) in horizontal_lines:
        x_coords.append(x1)
        x_coords.append(x2)
        y_coords.append(y)
    for x, (y1, y2) in vertical_lines:
        y_coords.append(y1)
        y_coords.append(y2)
        x_coords.append(x)
    x_map = get_compressed_coords(x_coords)
    y_map = get_compressed_coords(y_coords)

    comp_hor_lines = [(y_map[y], (x_map[x1], x_map[x2])) for y, (x1, x2) in horizontal_lines]
    comp_ver_lines = [(x_map[x], (y_map[y1], y_map[y2])) for x, (y1, y2) in vertical_lines]

    return comp_hor_lines, comp_ver_lines, len(x_map), len(y_map)


def get_merged_lines(lines):
    coord_lines = {}
    merged_lines = []

    for coord, (start, end) in lines:
        if coord not in coord_lines:
            coord_lines[coord] = []
        coord_lines[coord].append((start, end))

    for coord in coord_lines:
        sorted_coord_lines = sorted(coord_lines[coord], key=lambda x: x[0])
        curr_beg, curr_end = sorted_coord_lines[0]
        for next_beg, next_end in sorted_coord_lines[1:]:
            if next_beg <= curr_end:
                curr_end = max(curr_end, next_end)
            else:
                merged_lines.append((coord, (curr_beg, curr_end)))
                curr_beg, curr_end = next_beg, next_end
        merged_lines.append((coord, (curr_beg, curr_end)))

    return merged_lines