AtCoder Beginners Contest 331 (ABC331) をPythonで解きました。
見やすいコードを書く練習も兼ねてます。
A - Tomorrow
問題
$y$ 年 $m$ 月 $d$ 日の次の日は、何年何月何日ですか？
考察
次の日にちを求めたいので、とりあえず $d$ を $+1$ します。
1か月は $D$ 日までなので、もし $d$ が $D$ を超えていたら、 $d$ を $1$ にして、 $m$ を $+1$ します。
さらに、1年間は $M$ 日までなので、もし $m$ が $M$ を超えていたら、 $m$ を $1$ にして、 $y$ を $+1$ します。
コード
M, D = map(int, input().split())
y, m, d = map(int, input().split())
d += 1
if d > D:
d = 1
m += 1
if m > M:
m = 1
y += 1
print(y, m, d)
B - Buy One Carton of Milk
問題
$6$ 個入り、 $8$ 個入り、 $12$ 個入りの卵のパックがそれぞれ $S$, $M$, $L$ 円で売られています。卵を $N$ 個以上買うための最小の金額を求めてください。
考察
卵のパックは、どれも $N+1$ 個以上買う必要はありません。(本当はもっとこの上限の値は下げられますが、今回はそこまで計算しなくてもACできるのでこれでいきます)
制約を見てみると $N \leq 100$ なので、 卵のパックの数をそれぞれ $0$ ～ $N$ 個の中で全探索しても探索回数は $100^3=10^6$ 回程度です。
コード
# 全体の値段
def price(s, m, l):
return S * s + M * m + L * l
# 卵の個数
def cnt(s, m, l):
return 6 * s + 8 * m + 12 * l
N, S, M, L = map(int, input().split())
ans = 999999999
for s in range(N + 1):
for m in range(N + 1):
for l in range(N + 1):
if cnt(s, m, l) >= N:
ans = min(ans, price(s, m, l))
print(ans)
C - Sum of Numbers Greater Than Me
問題
数列 $A$ が与えられます。 $i=0,1,\cdots N$ について、次の問題を解いてください。
- $A$ の要素のうち、 $A_i$ よりも大きな要素すべての和を求めてください。
考察
愚直な解き方としては、各 $i=0,1,\cdots N$ について、 $A$ の中で $A_i$ よりも大きな数の和を出力する下のようなコードがあります。
"""愚直な解き方"""
N = int(input())
A = list(map(int, input().split()))
ans_lst = []
for i in range(N):
total = 0
for j in range(N):
if A[j] > A[i]: total += A[j]
ans_lst.append(total)
print(*ans_lst)
これは $O(N^2)$ なので、TLEしてしまいます。
なのでどうにか工夫したいです。
方針として、下のような辞書をつくりたいです。
"""辞書をつくりたい"""
N = int(input())
A = list(map(int, input().split()))
ans_dic = dict()
# ここにans_dicをつくる処理が入る。
ans_lst = []
for a in A:
ans_lst.append(ans_dic[a])
print(*ans_lst)
いろんな書き方がありますが、ここではsortedcontainersのSortedDictをつかったコードを紹介します。これをつかうと、かなりシンプルめに実装できます。
SortedDictは、普通の辞書と同じように扱えて、なおかつ常にキーが小さい順にソートされるような辞書です。
"""SortedDictの使用例"""
from sortedcontainers import SortedDict
dic = SortedDict()
dic[3] = "three"
dic[5] = "five"
dic[1] = "one"
for key, value in dic.items():
print(key, value)
# 出力結果
# 1 one
# 3 three
# 5 five
キーを $A$ の要素に、バリューをその個数にしたSortedDictを用意します。
これをつかうと、つくりたかった ans_dic をつくることができ、ACすることができます。
コード
from sortedcontainers import SortedDict
N = int(input())
A = list(map(int, input().split()))
count_dic = SortedDict()
for a in A:
if a not in count_dic:
count_dic[a] = 0
count_dic[a] += 1
total = 0
ans_dic = dict()
for key, count in reversed(count_dic.items()):
ans_dic[key] = total
total += key * count
print(*[ans_dic[a] for a in A])
D - Tile Pattern
問題
下のタイプの $Q$ 回のクエリに答えてください。
- 左上隅 $(A_i,B_i)$ 右下隅 $(C_i,D_i)$ の長方形の中にある黒マスの個数はいくつですか？
考察1 (方針)
1次元の問題で、 $[L,R)$ の範囲の何かを求めてくださいと問われたら、$f(x)=[0,R)の範囲の答え$ となるような $f(x)$ をつくって、$f(R)-f(L)$ を出力するのが典型です。
2次元でも同じで、左上隅が $(0,0)$ で、縦の長さが $r$ 横の長さが $c$ のときの黒マスの個数を求める関数 $f(r,c)$ をつくりたいです。
もしつくれたらどう解けるかを、先に考えてみましょう。
入力例1で、 $(A,B,C,D)=(1,2,3,4)$ のときの図です。
まず、$f(4,5)$ で、 $(0,0)$ から $(3,4)$ までの長方形の中にある黒マスの数を求めます。
無駄に数えてしまっている部分があるので、その部分を引き算したいです。
$f(4,2)$ と $f(1,5)$ を求めて引きます。
$f(1,2)$ の部分を2度引いてしまっているので、帳尻合わせで足しておきます。
まとめると、$(A,B,C,D)=(1,2,3,4)$ のとき、答えは $f(4,5)-f(4,2)-f(1,5)+f(1,2)$ です。
ここまでの解答コードはこんな感じになります。
"""考察1までのコード"""
# [0,r),[0,c) の範囲にあるBの数
def f(r, c):
pass # 考察2以降でやります。
N, Q = map(int, input().split())
P = [input() for _ in range(N)]
for _ in range(Q):
si, sj, gi, gj = map(int, input().split())
ans = f(gi + 1, gj + 1) - f(si, gj + 1) - f(gi + 1, sj) + f(si, sj)
考察2 (Pについての2次元累積和)
$N \leq 1000$ なので、$r,c \leq N$ のときは気合いで全探索すれば $O(N^2)$ で求められます。
しかし、これを各クエリごとにやるわけにはいかないので、$r,c \leq N$ のときの答えは先に計算してメモしておきましょう。
下のコードでは、いもす法を縦と横の両方にすることで求めています。
# 半開区間[0,i), [0,j)内のBの個数 (0<=i,j<=N)
P_acc = [[0] * (N + 1) for _ in range(N + 1)]
for i in range(N):
for j in range(N):
if P[i][j] == 'B':
P_acc[i + 1][j + 1] = 1
for i in range(1, N + 1):
for j in range(1, N + 1):
P_acc[i][j] += P_acc[i][j - 1]
for j in range(1, N + 1):
for i in range(1, N + 1):
P_acc[i][j] += P_acc[i - 1][j]
P_sum = P_acc[N][N] # P内のBの個数
考察3 (f(r,c)を求める)
最後に $f(r,c)$ を求めます。これが求まれば完成です。
制約より $r,c \leq 10^9$ なので、いちいち全部の黒マスの数を数えてられません。
$N \times N$ の形の $P$ がたくさんあるので、これを利用したいです。
割り算をして、 $r=q_i \times N + m_i$ , $c=q_j\times N + m_j$ とします。
オレンジの部分は、$N \times N$ の$P$ 内にある黒マスの数を sum とすると、 $sum \times q_i\times q_j$ です。この変数は、考察2で求めています。
緑・灰色の部分は、次のコードで求められます。
green = P_acc[mi][N] * qj
gray = P_acc[N][mj] * qi
青色の部分は P_acc をそのまま利用して求められます。
まとめると、解答コードは次のようになります。
コード
# 半開区間[0, r), [0, c) 内のBの個数
def f(r, c):
qi, mi = divmod(r, N)
qj, mj = divmod(c, N)
result = (qi * qj) * P_sum
result += P_acc[mi][N] * qj
result += P_acc[N][mj] * qi
result += P_acc[mi][mj]
return result
N, Q = map(int, input().split())
P = [input() for _ in range(N)]
# 半開区間[0,i), [0,j)内のBの個数 (0<=i,j<=N)
P_acc = [[0] * (N + 1) for _ in range(N + 1)]
for i in range(N):
for j in range(N):
if P[i][j] == 'B':
P_acc[i + 1][j + 1] = 1
for i in range(1, N + 1):
for j in range(1, N + 1):
P_acc[i][j] += P_acc[i][j - 1]
for j in range(1, N + 1):
for i in range(1, N + 1):
P_acc[i][j] += P_acc[i - 1][j]
P_sum = P_acc[N][N] # P内のBの個数
for _ in range(Q):
si, sj, gi, gj = map(int, input().split())
gi, gj = gi + 1, gj + 1
ans = f(gi, gj) - f(gi, sj) - f(si, gj) + f(si, sj)
print(ans)
E - Set Meal
問題
主菜と副菜の組み合わせのなかで、最も高い金額になるものはいくらですか？
ただし、食べ合わせの悪い組み合わせを選ぶのはナシです。
考察
副菜のリスト $B$ を大きい順にソートします。
主菜のリスト $A$ の要素 $a$ それぞれについて、ソートされたリスト $B$ の1つめの要素 $B_0$ と組み合わせて定食をつくります。
もし食べ合わせが悪いなら、$B_1$ と組み合わせ、食べ合わせが良くなるまでこれを繰り返します。
食べ合わせが良いとき、変数ansを更新します。
この $a$ について、さらに副菜リスト $B$ の要素を見る必要はありません。 $B$ はソートされていて、この先はより安い金額しか現れないからです。
これを $A$ の要素すべてについて行えばokです。
一見すると計算量が $O(NM+MlogM)$ に見えますが、食べ合わせの悪い組み合わせは $L$ 個なので、全体で $O(N+MlogM+L)$ です。
コード
N, M, L = map(int, input().split())
A = list(map(int, input().split()))
B = list(map(int, input().split()))
ng_set = set()
for _ in range(L):
c, d = map(lambda x: int(x) - 1, input().split())
ng_set.add((c, d))
p = sorted([(b, j) for j, b in enumerate(B)], reverse=True)
ans = -1
for i, a in enumerate(A):
for b, j in p:
if (i, j) not in ng_set:
ans = max(ans, a + b)
break
print(ans)