GigaCode2019D 家の建設

この記事で使うアルゴリズム

累積和

はじめに

カテゴリー競プロ初中級者100問では、Qiitaにて@e869120さんがレッドコーダーが教える、競プロ・AtCoder上達のガイドライン【中級編:目指せ水色コーダー!】としてまとめられている100問をPythonで解いています。

全問題の一覧はこちらです

問題

https://atcoder.jp/contests/gigacode-2019/tasks/gigacode_2019_d

この問題では累積和の計算が求められます。

方針

H方向の座標がh1+1,h1+2,...,h2、W方向の座標がw1+1,w1+2,...,w2である長方形についてマス内の和を求めたいとします。
この時、二次元の累積和Cを考えます。そして、C[h2][w2] - C[h1][w2] - C[h2][w1] + C[h1][w1]を求めれば良いです。

単純にこの解法だとPyPyであっても際どくTLEとなってしまったので、泥臭いですが一部の計算をメモして高速化しました。

コード

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
import itertools
def transpose(mat): # 二次元のリストを転置する関数
return list(zip(*mat))
H, W, K, V = map(int, input().split())
A = [[0] * (W+1)] # 1行目は0にしておくと計算が楽

for h in range(H):
a = list(map(int,input().split()))
a = [0] + a # 1列目は0にしておくと計算が楽
A.append(a)

B = []
for a in A:
csum = list(itertools.accumulate(a))
B.append(csum)
B = transpose(B)

csum_mat = []
for a in B:
csum = list(itertools.accumulate(a))
csum_mat.append(csum)
csum_mat = transpose(csum_mat)

ans = 0
for h1 in range(H):
if ((H-h1) * W < ans):
continue
ch1 = csum_mat[h1] # メモ
for h2 in range(h1,H+1):
if ((h2-h1) * W < ans):
continue
h2h1 = h2 - h1 # メモ
ch2 = csum_mat[h2] # メモ
for w1 in range(W):
if (h2h1 * (W-w1) < ans):
continue
h2w1h1w1 = - ch2[w1] + ch1[w1] # メモ
for w2 in range(w1,W+1):
s = h2h1 * (w2-w1) # 面積
if (s < ans):
continue
cost = ch2[w2] - ch1[w2] + h2w1h1w1
cost += K * s
if (cost <= V):
ans = s
print (ans)

記事情報

  • 投稿日:2020年5月5日
  • 最終更新日:2020年5月8日