|
| 1 | +#!/usr/bin/env python3 |
| 2 | +# -*- coding: utf-8 -*- |
| 3 | + |
| 4 | +import array |
| 5 | +import collections |
| 6 | +import heapq |
| 7 | + |
| 8 | + |
| 9 | +AdjacentVertex = collections.namedtuple("AdjacentVertex", "vertex cost") |
| 10 | +INF = 2 ** 31 - 1 |
| 11 | +NO_VERTEX = -1 |
| 12 | + |
| 13 | + |
| 14 | +# Prim法で頂点0からの最小全域木を求める |
| 15 | +def compute_mst_prim(max_v, adj_list): |
| 16 | + # pred[u]は頂点uの「ひとつ前」の頂点を表す |
| 17 | + pred = collections.defaultdict(lambda: NO_VERTEX) |
| 18 | + # uとpred[u]を結ぶ辺の重みがkey[u]に入る |
| 19 | + key = collections.defaultdict(lambda: INF) |
| 20 | + key[0] = 0 |
| 21 | + # 二分ヒープを優先度付きキューとして用いる |
| 22 | + pq = [(key[v], v) for v in range(max_v)] |
| 23 | + heapq.heapify(pq) |
| 24 | + # 優先度付きキューに頂点が入っているかを示す配列 |
| 25 | + in_pq = array.array("B", (True for _ in range(max_v))) |
| 26 | + while pq: |
| 27 | + _, u = heapq.heappop(pq) |
| 28 | + in_pq[u] = False |
| 29 | + for v, v_cost in adj_list[u]: |
| 30 | + if in_pq[v]: |
| 31 | + weight = v_cost |
| 32 | + if weight < key[v]: |
| 33 | + pred[v] = u |
| 34 | + key[v] = weight |
| 35 | + heapq.heappush(pq, (weight, v)) |
| 36 | + in_pq[v] = True |
| 37 | + return (pred, key) |
| 38 | + |
| 39 | + |
| 40 | +def main(): |
| 41 | + W, H = list(map(int, input().split())) |
| 42 | + # max_v, max_e = map(int, input().split()) |
| 43 | + max_v = (W+1) * (H+1) |
| 44 | + max_e = (W+1)*H + (H+1)*W |
| 45 | + adjacency_list = collections.defaultdict(set) |
| 46 | + for _w in range(W): |
| 47 | + p = int(input()) |
| 48 | + for _h in range(H): |
| 49 | + adjacency_list[_w*H+_h].add(AdjacentVertex((_w+1)*H+_h, p)) |
| 50 | + adjacency_list[(_w+1)*H+_h].add(AdjacentVertex(_w*H+_h, p)) |
| 51 | + for _h in range(H): |
| 52 | + q = int(input()) |
| 53 | + for _w in range(W): |
| 54 | + adjacency_list[_w*H+_h].add(AdjacentVertex((_w+1)*H+_h, q)) |
| 55 | + adjacency_list[(_w+1)*H+_h].add(AdjacentVertex(_w*H+_h, q)) |
| 56 | + # for _ in range(max_e): |
| 57 | + # s, t, w = map(int, input().split()) |
| 58 | + # adjacency_list[s].add(AdjacentVertex(t, w)) |
| 59 | + # adjacency_list[t].add(AdjacentVertex(s, w)) |
| 60 | + (_, key) = compute_mst_prim(max_v, adjacency_list) |
| 61 | + print(sum(key.values())) |
| 62 | + |
| 63 | + |
| 64 | +if __name__ == '__main__': |
| 65 | + main() |
| 66 | + |
0 commit comments