競プロ例題【ちょいムズ】【★４】【AtCoder】【茶色】

だいぶ雑ですが、一旦アップします。

問題
1. 問題文
解説
1. 大枠
2. 中身（１番大事）
  1. 計算量改善

問題

012 - Red Painting（★4）

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

こちらの問題の解説記事です。（もちろん非公式です。）

リンク先の問題を１度解こうとしてから来てもらうと本記事も見やすいかと思います。

問題文

H 行 W 列のマス目があり、上から i (1≤i≤H) 行目、左から j (1≤j≤W) 列目のマスを (i,j) と表します。
最初すべてのマスは白いです。ここで、Q 個のクエリが以下の形式で与えられます。
i (1≤i≤Q) 番目のクエリについて：
ti=1 のとき
整数 ri, ci が与えられる。
白いマス (ri,ci) が赤色で塗られる。
ti=2 のとき
整数 rai, cai, rbi, cbi が与えられる。
次の 2 つの条件両方を満たすとき Yes、そうでなければ No と出力する。
マス (rai,cai) とマス (rbi,cbi) が赤色で塗られている。
マス (rai,cai) からマス (rbi,cbi) まで赤色マス上を上下左右に移動することで辿り着ける。
以上の Q 個のクエリを順に処理してください。

012 - Red Painting（★4）

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

解説

以下、めっちゃ自己流ですが、つらつらと書いていきます。

大枠

まず、思考の整理のために大枠を書きました。

def hantei:

H, W = map(int, input().split())
Q = int(input())
for i in range(Q):
    q = list(map(int, input().split()))

#マスを作る
masu = [[0]*W for i in range(H)]

for i in range(Q):
    if q[i] == 1:
        masu[q[i][1]][q[i][2]] = 1 #赤色に変更
    else:
        if masu[q[i][1]][q[i][2]] == 1 and masu[q[i][3]][q[i][4]] == 1:
            if hantei:#繋がっているかの判定
                print("Yes")
            else:
                print("No")
        else:
            print("No")

「if hantei:#繋がっているかの判定」のhantei関数さえ作れたらおっけいと言う状態にしました。そしてここが肝なわけですが、一旦全体を把握というか、入出力を明確にして何に集中すればいいかを狭めたかったので、こんな感じで大枠を書いてみました。

中身（１番大事）

変数名とかも適宜変えてます。

#中身
from collections import deque

def bfs(maze, start, end): #startからendまで繋がっているかどうかを返す
    H, W = len(maze), len(maze[0])
    queue = deque([(start[0], start[1])])
    visited = [[False] * W for _ in range(H)]
    visited[start[0]][start[1]] = True

    while queue:
        i, j = queue.popleft()
        if (i, j) == end:
            return True

        for di, dj in [(0, 1), (0, -1), (1, 0), (-1, 0)]:#１マス移動
            ni, nj = i + di, j + dj
            if 0 <= ni < H and 0 <= nj < W and maze[ni][nj] == 1 and not visited[ni][nj]:
                visited[ni][nj] = True
                queue.append((ni, nj))

    return False

#全体
from collections import deque

def bfs(maze, start, end): #startからendまで繋がっているかどうかを返す
    H, W = len(maze), len(maze[0])
    queue = deque([(start[0], start[1])])
    visited = [[False] * W for _ in range(H)]
    visited[start[0]][start[1]] = True

    while queue:
        i, j = queue.popleft()
        if (i, j) == end:
            return True

        for di, dj in [(0, 1), (0, -1), (1, 0), (-1, 0)]:#１マス移動
            ni, nj = i + di, j + dj
            if 0 <= ni < H and 0 <= nj < W and maze[ni][nj] == 1 and not visited[ni][nj]:
                visited[ni][nj] = True
                queue.append((ni, nj))

    return False

H, W = map(int, input().split())
Q = int(input())
maze = [[0] * W for _ in range(H)]

for _ in range(Q):
    query = list(map(int, input().split()))
    if query[0] == 1:
        maze[query[1] - 1][query[2] - 1] = 1
    else:
        start = (query[1] - 1, query[2] - 1)
        end = (query[3] - 1, query[4] - 1)
        if maze[start[0]][start[1]] == 1 and maze[end[0]][end[1]] == 1:
            if bfs(maze, start, end):
                print("Yes")
            else:
                print("No")
        else:
            print("No")

bfs（Breadth-first search/ 幅優先探索）
https://qiita.com/drken/items/996d80bcae64649a6580

計算量改善

さっきのままだと３重ループなので計算を軽くします。

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.size = [1] * n

    def find(self, x):
        if self.parent[x] == x:
            return x
        else:
            self.parent[x] = self.find(self.parent[x])
            return self.parent[x]

    def union(self, x, y):
        x_root = self.find(x)
        y_root = self.find(y)
        if x_root != y_root:
            if self.size[x_root] < self.size[y_root]:
                x_root, y_root = y_root, x_root
            self.parent[y_root] = x_root
            self.size[x_root] += self.size[y_root]

H, W = map(int, input().split())
Q = int(input())
maze = [[0] * W for _ in range(H)]
uf = UnionFind(H * W)

for _ in range(Q):
    query = list(map(int, input().split()))
    if query[0] == 1:
        i, j = query[1] - 1, query[2] - 1
        maze[i][j] = 1
        for di, dj in [(0, 1), (0, -1), (1, 0), (-1, 0)]:
            ni, nj = i + di, j + dj
            if 0 <= ni < H and 0 <= nj < W and maze[ni][nj] == 1:
                uf.union(i * W + j, ni * W + nj)
    else:
        start = (query[1] - 1, query[2] - 1)
        end = (query[3] - 1, query[4] - 1)
        if maze[start[0]][start[1]] == 1 and maze[end[0]][end[1]] == 1:
            if uf.find(start[0] * W + start[1]) == uf.find(end[0] * W + end[1]):
                print("Yes")
            else:
                print("No")
        else:
            print("No")