Pythonの新常識 dataclassを使ってみた。AtCoder 二次元問題用処理を書いてみる編

はじめに

dataclassについて勉強中です。
dataclass自体の説明は、他の記事にお任せします。

学習の為にdataclassを使ってみたかったので、AtCoderの問題で使ってみました。

2次元問題

私が勝手に2次元問題と呼んでいるのですが、H列、W行の二次元平面の問題があります。
例えばこの問題

(1,1) から(H,W) まで移動できるは判定する問題です。
上下左右に移動する時、[(1,0),(-1,0),(0,1),(0,-1)]のようなリストを用意して移動させたり、H列、W行の範囲外のチェックなどが必要になります。

dfsを使って辿っていけるかを確認してみます。
二次元固有の処理はBoardクラスにまとめてみます。

from dataclasses import dataclass

@dataclass
class Board:
    H: int
    W: int
    S: list

    def get_neighbor_pos(self, pos):
        i,j = pos
        for di,dj in [(1,0),(-1,0),(0,1),(0,-1)]:
            ni = i + di
            nj = j + dj
            if 0<= ni < self.H and 0 <= nj < self.W:
                yield (ni,nj)
    def get(self, pos):
        i,j = pos
        return self.S[i][j]

def solve(H: int, W: int, S: "List[str]"):
    
    s = ["s","n","u","k","e"]
    visited = defaultdict(bool)

    board = Board(H,W,S)

    def dfs(pos,d):
        if pos == (H-1,W-1):
            return True
        for next_pos in board.get_neighbor_pos(pos):
            if not visited[next_pos] and board.get(next_pos) == s[(d+1)%5]:
                visited[next_pos] = True
                if dfs(next_pos, d+1):
                    return True
        return False

    if dfs((0,0),0):
        print(YES)
    else:
        print(NO)

まとめ

• classを使った記述では__init__()の記述に比べてすっきり書ける。
• dfsの関数中身は、経路をたどっていく処理に専念した記述にできた。
• このclassをもちょっと育てていったら、使いやすいものになりそう。

AtCoderの問題は、いつも力づくで書いてしまうんですが、綺麗な構造で書きたいですね。