• Skynet0's avatar
    Code · 9f5fe1b7
    Skynet0 authored
    9f5fe1b7
country_util.py 1.54 KB
from typing import Dict, List, NamedTuple, Tuple

from cs11puzzles.pzpr_helpers import conns_to_str_grid


class CountryInstance(NamedTuple):
    areas: List[List[int]]
    area_nums: Dict[int, int]


def load_puzzle_and_solution(fname: str) -> Tuple[CountryInstance, List[List[str]]]:
    with open(fname) as f:
        assert f.readline().strip() == "pzprv3.1"
        assert f.readline().strip() == "country"

        height = int(f.readline())
        width = int(f.readline())
        f.readline()  # empty line for some reason
        _ = int(f.readline())  # n_regions

        areas = [
            [int(n) for n in f.readline().strip().split(" ")] for _ in range(height)
        ]

        numbers = [
            [int(n) if n != "." else 0 for n in f.readline().strip().split(" ")]
            for _ in range(height)
        ]
        area_nums = {}
        # pzpr displays number in leftmost, topmost cell
        for x in range(width):
            for y in range(height):
                if areas[y][x] not in area_nums:
                    area_nums[areas[y][x]] = numbers[y][x]

        conn_right = [
            [int(n) == 1 for n in f.readline().strip().split(" ")]
            for _ in range(height)
        ]
        conn_down = [
            [int(n) == 1 for n in f.readline().strip().split(" ")]
            for _ in range(height - 1)
        ]

        return CountryInstance(areas, area_nums), conns_to_str_grid(
            conn_right, conn_down
        )


def load_puzzle(fname: str):
    puzzle, _ = load_puzzle_and_solution(fname)
    return puzzle