#1195076 (PyPy3) No.5025 Forest Fire

• 問題: No.5025 Forest Fire
• ユーザー: prussian_coder
• 提出日時: 2026-10-24 16:41:21
• 言語: PyPy3 (8.0.0 + ACL)
• 結果: AC
• 実行時間: 1,826 ms / 2,000 ms
• コード長: 4,755 bytes
• コンパイル時間: 81 ms
• コンパイル使用メモリ: 64,356 KB
• 実行使用メモリ: 95,432 KB
• スコア: 161,799,450
• 最終ジャッジ日時: 2026-10-24 16:44:44
• 合計ジャッジ時間: 148,768 ms
• ジャッジサーバーID （参考情報）: judge4_0 / <nil>

ファイルパターン	結果
other	AC * 50

ソースコード

import random
import time
start_time = time.time()
N,K=map(int,input().split())
A = [[int(x) for x in input()] for _ in range(N)]
N2 = N+2

BOARD = [[0]*N2 for _ in range(N2)]

for i in range(N):
    for j in range(N):
        BOARD[i+1][j+1]=A[i][j]

SAFE_ZONE_CRITERIA = 5
CAND_SIZES_LIM = 1000

def to_1d(i,j):
    return i*N2+j

def to_2d(v):
    return v//N2, v%N2

def connection_dist(v1,v2):
    x1,y1 = to_2d(v1)
    x2,y2 = to_2d(v2)
    return min(abs(x1-x2),abs(y1-y2))

def create_route(v1,v2):
    x1,y1 = to_2d(v1)
    x2,y2 = to_2d(v2)
    if x1>x2:
        x1,x2 = x2,x1
        y1,y2 = y2,y1
    route = []
    x,y = x1,y1
    while (x,y)!=(x2,y2):
        if x<x2:
            x+=1
        if y<y2:
            y+=1
        elif y>y2:
            y-=1
        if (x,y)!=(x2,y2):
            route.append(to_1d(x,y))
    return route
        
from collections import defaultdict

class UnionFind():
    def __init__(self, n):
        self.n = n
        self.parents = [-1] * n

    def find(self, x):
        if self.parents[x] < 0:
            return x
        else:
            self.parents[x] = self.find(self.parents[x])
            return self.parents[x]

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)

        if x == y:
            return

        if self.parents[x] > self.parents[y]:
            x, y = y, x

        self.parents[x] += self.parents[y]
        self.parents[y] = x

    def size(self, x):
        return -self.parents[self.find(x)]

    def same(self, x, y):
        return self.find(x) == self.find(y)

    def members(self, x):
        root = self.find(x)
        return [i for i in range(self.n) if self.find(i) == root]

    def roots(self):
        return [i for i, x in enumerate(self.parents) if x < 0]

    def group_count(self):
        return len(self.roots())

    def all_group_members(self):
        group_members = defaultdict(set)
        for member in range(self.n):
            group_members[self.find(member)].add(member)
        return group_members

    def __str__(self):
        return '\n'.join(f'{r}: {m}' for r, m in self.all_group_members().items())

safe_zones = []
danger_zones = set()
for i in range(N2):
    for j in range(N2):
        if BOARD[i][j]<=SAFE_ZONE_CRITERIA:
            safe_zones.append(to_1d(i,j))
        else:
            danger_zones.add(to_1d(i,j))


edge_candidates = set()

while len(edge_candidates)<=CAND_SIZES_LIM:
    v1 = random.choice(safe_zones)
    v2 = random.choice(safe_zones)
    if 2<=connection_dist(v1,v2)<=12:
        edge_candidates.add((v1,v2))

edge_candidates = list(edge_candidates)

def build_trees(connections: list[tuple[int,int]]):
    tree_count = 0
    trees = set()
    used_connections = []
    for v1,v2 in connections:
        route = create_route(v1,v2)
        for v in route:
            if v in trees:
                continue
            x,y = to_2d(v)
            if BOARD[x][y]==0:
                continue
            trees.add(v)
            tree_count+=1
            if tree_count==K:
                return trees
    return trees

def evalulate(trees):
    targets = danger_zones - trees
    m = len(targets)
    targets_idx = {v:i for i,v in enumerate(targets)}
    uf = UnionFind(m)
    for v in targets:
        for d in [1,-1,N2,-N2]:
            if v+d not in targets:
                continue
            uf.union(targets_idx[v], targets_idx[v+d])
    score = 0
    for u in uf.all_group_members().values():
        score += len(u)**2
    return score

best_score = float("inf")
best_trees = None

connections = []
est_sizes = 0
edges = []
current_score = float("inf")
best_score = float("inf")

while time.time()-start_time <= 1.8:
    x = random.randint(0, CAND_SIZES_LIM)
    r = random.random()
    if x in edges:
        continue
    if est_sizes <= 90:
        v1,v2 = edge_candidates[x]
        connections.append((v1,v2))
        est_sizes += connection_dist(v1,v2)
        edges.append(x)
        prev_edge = None
    else:
        y = random.randint(0, len(edges)-1)
        prev_edge = edges[y]
        prev_connection = connections[y]
        prev_est_sizes = est_sizes
        edges[y]=x
        connections[y]=edge_candidates[x]
        est_sizes -= connection_dist(*prev_connection)
        est_sizes += connection_dist(*edge_candidates[x])
    trees = build_trees(connections)
    new_score = evalulate(trees)
    if new_score <= current_score:
        current_score = new_score
    else:
        edges[y]=prev_edge
        connections[y]=prev_connection
        est_sizes = prev_est_sizes
    if new_score < best_score:
        best_trees = list(trees)
        best_score = new_score
print(len(best_trees))
for v in best_trees:
    x,y = to_2d(v)
    print(x-1,y-1)