advent-of-code/2024/python/matrix.py

from copy import deepcopy
from collections import defaultdict
import math

from typing import List, Dict, Tuple

split_word_to_chr_list = lambda y: [w for w in y]
split_word_to_int_list = lambda y: [int(w) for w in y]
split_line_to_int_list = lambda y: [int(w) for w in y.split(" ") if w]


def split_x_out(l):
    return [x for _, x in l]


def split_y_out(l):
    return [y for y, _ in l]


class colors:
    # HEADER = '\033[95m'
    BLUE = "\033[94m"
    GREEN = "\033[92m"
    YELLOW = "\033[93m"
    RED = "\033[91m"
    ENDC = "\033[0m"
    BLINK = "\033[5m"


def apply_to_all(mx, func):
    for row_num, row in enumerate(mx):
        for col_num, val in enumerate(row):
            mx[row_num][col_num] = func(val)


def rotate(m, right=True):  # -90
    """
    Takes a matrix, and rotates all of the values 90 degrees to the left
    """
    x = list(zip(*m[::-1]))
    if right:
        return x
    return [list(reversed(y)) for y in x]


def load_matrix_file(name, func=None):
    """
    Open a file and split all space separated word lists to integers as a matrix
    """
    with open(name, "r") as f:
        my_file = []
        for line in f:
            my_file.append(line.rstrip())
    if func:
        return [func(x) for x in my_file]
    return [split_word_to_int_list(x) for x in my_file]


def number_starting_at(mx, x, y):
    current = ""
    try:
        d = mx[y][x]
    except IndexError:
        return None

    if not d.isdigit():
        return None

    current += d
    next = number_starting_at(mx, x + 1, y)
    if next is None:
        return current
    return current + next


def find_in_matrix(mx, what, one=True):
    coords = []
    for row_num, row in enumerate(mx):
        for col_num, val in enumerate(row):
            if val == what:
                coord = (row_num, col_num)
                if one is True:
                    return coord
                else:
                    coords.append(coord)
    return coords


def get_neighbors(matrix, x, y, _dict=False):
    neighbors = []
    # left
    try:
        if x - 1 >= 0:
            if _dict:
                neighbors.append({"x": x - 1, "y": y, "value": matrix[y][x - 1]})
            else:
                neighbors.append([(x - 1, y), matrix[y][x - 1]])
    except IndexError:
        pass
    # right
    try:
        if _dict:
            neighbors.append({"x": x + 1, "y": y, "value": matrix[y][x + 1]})
        else:
            neighbors.append([(x + 1, y), matrix[y][x + 1]])
    except IndexError:
        pass

    # up
    try:
        if y - 1 >= 0:
            if _dict:
                neighbors.append({"x": x, "y": y - 1, "value": matrix[y - 1][x]})
            else:
                neighbors.append([(x, y - 1), matrix[y - 1][x]])
    except IndexError:
        pass

    # down
    try:
        if _dict:
            neighbors.append({"x": x, "y": y + 1, "value": matrix[y + 1][x]})
        else:
            neighbors.append([(x, y + 1), matrix[y + 1][x]])
    except IndexError:
        pass

    return neighbors


def valid_neighbors(matrix, x, y, criteria=None):
    if criteria is None:
        raise Exception("Please pass in a lambda for criteria")
    cur = matrix[y][x]
    neighbors = get_neighbors(matrix, x, y, _dict=True)
    valid = []
    for neighbor in neighbors:
        if criteria(cur, neighbor["value"]):
            valid.append(neighbor)
    return valid


def sum_matrix(mtx):
    total = 0
    for row in mtx:
        total += sum(row)
    return total


M_UL, M_U, M_UR = (-1, -1), (0, -1), (1, -1)
M_L, M_R = (-1, 0), (1, 0)
M_DL, M_D, M_DR = (-1, 1), (0, 1), (1, 1)

M_NW, M_N, M_NE = (-1, -1), (0, -1), (1, -1)
M_W, M_E = (-1, 0), (1, 0)
M_SW, M_S, M_SE = (-1, 1), (0, 1), (1, 1)
CARDINALS = {
    M_NW: "NW",
    M_N: "N",
    M_NE: "NE",
    M_W: "W",
    M_E: "E",
    M_SW: "SW",
    M_S: "S",
    M_SE: "SE",
}


def get_neighbor_coords(matrix, c, r, diagonals=True):
    height = len(matrix)
    width = len(matrix[0])
    if diagonals:
        coords = (M_UL, M_U, M_UR, M_L, M_R, M_DL, M_D, M_DR)
    else:
        coords = (M_U, M_L, M_R, M_D)
    neighbors = []

    for _c, _r in coords:
        try:
            value = matrix[r + _r][c + _c]  # Try to get a value error
            if r + _r >= 0 and c + _c >= 0:
                neighbors.append(
                    [{"c": c + _c, "r": r + _r}, value]
                )  # woo, no error, this coord is valid
        except IndexError:
            pass  # okay we out of bounds boizzzz
    return neighbors


def get_neighbors_cardinal(matrix, c, r, diagonals=True):
    height = len(matrix)
    width = len(matrix[0])
    coords = [M_N, M_W, M_E, M_S]
    if diagonals:
        coords.extend([M_NW, M_NE, M_SW, M_SE])
    neighbors = {}

    for coord in coords:
        _c, _r = coord
        try:
            value = matrix[r + _r][c + _c]  # Try to get a value error
            if r + _r >= 0 and c + _c >= 0:
                neighbors[CARDINALS[coord]] = {
                    "c": c + _c,
                    "r": r + _r,
                    "value": value,
                }  # woo, no error, this coord is valid
        except IndexError:
            pass  # okay we out of bounds boizzzz
    return neighbors


def line_of_sight_coords_diagonals(
    matrix, row, col, distance=None
) -> Dict[str, List[Tuple[int, int]]]:
    """
    Takes a matrix, a row, and a column
    calculates the coordinates to the edge for all four cardinal directions

    returns a dict with a list  of tuple coordes TRAVELING AWAY from the
    requested coordinate
    """
    height, width = get_size(matrix)

    col_ids = list(range(0, height))
    row_ids = list(range(0, width))

    directions = {
        "UR": [],
        "UL": [],
        "DR": [],
        "DL": []
    }
    ur = False
    ul = False
    dl = False
    dr = False
    
    # Boundaries of the grid
    max_row, max_col = height, width
    
    # Up-right (decrease row, increase column)
    r, c = row, col
    while r > 0 and c < max_col - 1:
        ur = True
        r -= 1
        c += 1
        directions["UR"].append((r, c))
    
    # Up-left (decrease row, decrease column)
    r, c = row, col
    while r > 0 and c > 0:
        ul = True
        r -= 1
        c -= 1
        directions["UL"].append((r, c))
    
    # Down-right (increase row, increase column)
    r, c = row, col
    while r < max_row - 1 and c < max_col - 1:
        dr = True
        r += 1
        c += 1
        directions["DR"].append((r, c))
    
    # Down-left (increase row, decrease column)
    r, c = row, col
    while r < max_row - 1 and c > 0:
        dl = True
        r += 1
        c -= 1
        directions["DL"].append((r, c))

    if distance:
        for _dir in ["DL", "UL", "DR", "UR"]:
            directions[_dir] =  directions[_dir][:distance - 1]
            if directions[_dir]:
                directions[_dir].insert(0, (row,col))
    return directions


def line_of_sight_coords(
    matrix, row, col, distance=None, diagonals=False
) -> Dict[str, List[Tuple[int, int]]]:
    """
    Takes a matrix, a row, and a column
    calculates the coordinates to the edge for all four cardinal directions

    returns a dict with a list  of tuple coordes TRAVELING AWAY from the
    requested coordinate
    """
    height, width = get_size(matrix)

    col_ids = list(range(0, height))
    row_ids = list(range(0, width))

    up_ids = list(reversed(col_ids[:col]))
    down_ids = col_ids[col + 1 :]
    left_ids = list(reversed(row_ids[:row]))
    right_ids = row_ids[row + 1 :]

    if distance:
        left = [(r, col) for r in left_ids[:distance - 1]]
        right = [(r, col) for r in right_ids[:distance - 1]]
        up = [(row, c) for c in up_ids[:distance - 1]]
        down = [(row, c) for c in down_ids[:distance - 1]]
        left.insert(0, (row, col))
        right.insert(0, (row, col))
        down.insert(0, (row, col))
        up.insert(0, (row, col))
    else:
        left = [(r, col) for r in left_ids]
        right = [(r, col) for r in right_ids]
        up = [(row, c) for c in up_ids]
        down = [(row, c) for c in down_ids]

    directions = {
        "U": up,
        "L": left,
        "D": down,
        "R": right,
    }
    if diagonals:
        vals = line_of_sight_coords_diagonals(matrix, row, col, distance)
        directions.update(vals)
    return directions


def line_of_sight(mx, row, col, distance=None, diagonals=False):
    """
    renders a line of sight coord calculation, into the values
    """
    coords = line_of_sight_coords(mx, row, col, distance, diagonals)
    los = defaultdict(list)
    for k, ids in coords.items():
        for _row, _col in ids:
            los[k].append(mx[_row][_col])
    return los


def coords_between_points(point1, point2):
    y1, x1 = point1
    y2, x2 = point2

    coords = []
    x = 0
    y = 0

    if x2 < x1:
        y = point1[0]
        for _x in range(x2, x1 + 1):
            coords.append((y, _x))
    elif x1 < x2:
        y = point1[0]
        for _x in range(x1, x2 + 1):
            coords.append((y, _x))
    elif y2 < y1:
        x = point1[1]
        for _y in range(y2, y1 + 1):
            coords.append((_y, x))
    elif y1 < y2:
        x = point1[1]
        for _y in range(y1, y2 + 1):
            coords.append((_y, x))
    return coords


def get_size(matrix):
    height = len(matrix)
    width = len(matrix[0])
    return height, width


def row_col_from_int(matrix, x):
    h, w = get_size(matrix)
    col = x % w
    row = x // h
    return row, col


def matrix_of_size(width, height, default=0):
    return [[default] * width for x in range(height)]


def set_matrix_dict(m):
    for x in range(len(m)):
        for y in range(len(m[x])):
            m[x][y] = {}
    return m


def pmx(*matrices, pad=True, space=True):
    """
    print a matrix of integers, zero turns to `.` for clarity
    """
    if len(matrices) > 1:
        matrices = list(zip(*matrices))
        for row in matrices:
            r = []
            for col in row:
                r.append("".join([f"{int(x)or '.'}".rjust(3) for x in col]))
            print(" ".join(r))
    else:
        for row in matrices:
            for c in row:
                if pad:
                    f = lambda x: f"{int(x)or '.'}".rjust(2)
                    if space:
                        f = lambda x: f"{int(x)or '.'}".rjust(3)
                else:
                    f = lambda x: f"{int(x)or '.'}"
                    if space:
                        f = lambda x: f"{int(x)or '.'} "
                print("".join([f(x) for x in c]))


def ppmx(*matrices, pad=True, space=True, zero="."):
    """
    print a matrix of anything, Falsy values turns to `.` for clarity
    """
    out = []
    if len(matrices) > 1:
        matrices = list(zip(*matrices))
        for row in matrices:
            r = []
            for col in row:
                r.append("".join([f"{x or zero}".rjust(3) for x in col]))
            out.append(" ".join(r))
    else:
        for row in matrices:
            for c in row:
                if pad:
                    f = lambda x: f"{x or zero}".rjust(2)
                    if space:
                        f = lambda x: f"{x or zero}".rjust(3)
                else:
                    f = lambda x: f"{x or zero}"
                    if space:
                        f = lambda x: f"{x or zero} "
                out.append("".join([f(x) for x in c]))
    return "\n".join(out)


def view_matrix(matrix, y1, x1, y2, x2):
    lines = ppmx(matrix, pad=0, space=0).split("\n")
    for line in lines[y1 : y2 + 1]:
        print(line[x1:x2])


def highlight(matrix, red=[], green=[], blue=[], blink_green=[]):
    """
    print a matrix of anything, Falsy values turns to `.` for clarity
    """
    mx = deepcopy(matrix)
    for y, x in red:
        if (y, x) in blue or (y, x) in green or (y, x) in blink_green:
            continue
        new = f"{colors.RED}{mx[y][x]}{colors.ENDC}"
        mx[y][x] = new
    for y, x in green:
        if (y, x) in blue or (y, x) in blink_green:
            continue
        new = f"{colors.GREEN}{mx[y][x]}{colors.ENDC}"
        mx[y][x] = new
    for y, x in blue:
        if (y, x) in blink_green:
            continue
        new = f"{colors.BLUE}{mx[y][x]}{colors.ENDC}"
        mx[y][x] = new
    for y, x in blink_green:
        new = f"{colors.BLINK}{colors.GREEN}{mx[y][x]}{colors.ENDC}"
        mx[y][x] = new
    print(ppmx(mx, pad=False, space=True, zero="0"))


def draw_shape_at(mx, row, col, shape=None, value=1):
    if shape is None:
        raise Exception("Please provide a list of coordinate offsets from Y,X to draw")
    for y, x in shape:
        mx[row + y][col + x] = value


def collision_at(mx, row, col, shape=None):
    if shape is None:
        raise Exception("Please provide a list of coordinate offsets from Y,X to draw")
    for y, x in shape:
        if mx[row + y][col + x] != 0:
            return True
    return False


def out_of_bounds(mx, row, col, shape=None):
    if shape is None:
        raise Exception("Please provide a list of coordinate offsets from Y,X to draw")
    height, width = get_size(mx)
    for y, x in shape:
        if row + y > height - 1:
            return True
        if col + x >= width:
            return True
        if col + x < 0:
            return True
    return False


def spiral_generator(width, height):
    k = 0
    l = 0
    m = height
    n = width

    """ k - starting row index
        m - ending row index
        l - starting column index
        n - ending column index
        i - iterator """

    while k < m and l < n:
        # Print the first row from
        # the remaining rows
        for i in range(l, n):
            yield (i, k)
            # print(a[k][i], end=" ")

        k += 1

        # Print the last column from
        # the remaining columns
        for i in range(k, m):
            yield (n - 1, i)
            # print(a[i][n - 1], end=" ")

        n -= 1

        # Print the last row from
        # the remaining rows
        if k < m:
            for i in range(n - 1, (l - 1), -1):
                # print(a[m - 1][i], end=" ")
                yield (i, m - 1)

            m -= 1

        # Print the first column from
        # the remaining columns
        if l < n:
            for i in range(m - 1, k - 1, -1):
                # print(a[i][l], end=" ")
                yield (l, i)

            l += 1