Kacemath's picture
feat: update with latest changes
47bba68
"""Grid and Segment models for the delivery search problem."""
from dataclasses import dataclass, field
from typing import Dict, Tuple, Optional
@dataclass
class Segment:
"""Represents a road segment between two adjacent grid points."""
src: Tuple[int, int]
dst: Tuple[int, int]
traffic: int # 0 = blocked, 1-4 = traffic level
def __post_init__(self):
# Normalize segment direction (ensure src < dst lexicographically)
if self.src > self.dst:
self.src, self.dst = self.dst, self.src
@property
def is_blocked(self) -> bool:
return self.traffic == 0
def get_key(self) -> Tuple[Tuple[int, int], Tuple[int, int]]:
"""Get normalized key for segment lookup."""
return (self.src, self.dst)
@dataclass
class Grid:
"""Represents the city grid with all road segments."""
width: int
height: int
segments: Dict[Tuple[Tuple[int, int], Tuple[int, int]], Segment] = field(
default_factory=dict
)
def get_segment(
self, src: Tuple[int, int], dst: Tuple[int, int]
) -> Optional[Segment]:
"""Get segment between two points (order doesn't matter)."""
key = (src, dst) if src < dst else (dst, src)
return self.segments.get(key)
def get_traffic(self, src: Tuple[int, int], dst: Tuple[int, int]) -> int:
"""Get traffic level for segment between two points."""
segment = self.get_segment(src, dst)
return segment.traffic if segment else 0
def is_blocked(self, src: Tuple[int, int], dst: Tuple[int, int]) -> bool:
"""Check if segment between two points is blocked."""
return self.get_traffic(src, dst) == 0
def is_valid_position(self, pos: Tuple[int, int]) -> bool:
"""Check if position is within grid bounds."""
x, y = pos
return 0 <= x < self.width and 0 <= y < self.height
def add_segment(self, src: Tuple[int, int], dst: Tuple[int, int], traffic: int):
"""Add or update a segment."""
segment = Segment(src, dst, traffic)
self.segments[segment.get_key()] = segment
def get_neighbors(self, pos: Tuple[int, int]) -> list[Tuple[int, int]]:
"""Get all valid neighboring positions (not blocked)."""
x, y = pos
neighbors = []
directions = [(0, 1), (0, -1), (1, 0), (-1, 0)] # up, down, right, left
for dx, dy in directions:
new_pos = (x + dx, y + dy)
if self.is_valid_position(new_pos) and not self.is_blocked(pos, new_pos):
neighbors.append(new_pos)
return neighbors
def to_dict(self) -> dict:
"""Convert grid to dictionary for JSON serialization."""
return {
"width": self.width,
"height": self.height,
"segments": [
{
"src": {"x": seg.src[0], "y": seg.src[1]},
"dst": {"x": seg.dst[0], "y": seg.dst[1]},
"traffic": seg.traffic,
}
for seg in self.segments.values()
],
}