ABC331をPythonで解いてみたよ。(A～E問題)

[AtCoder Beginners Contest 331 (ABC331)](https://atcoder.jp/contests/abc331) をPythonで解きました。
見やすいコードを書く練習も兼ねてます。


# A - Tomorrow

## 問題
$y$ 年 $m$ 月 $d$ 日の次の日は、何年何月何日ですか？

https://atcoder.jp/contests/abc331/tasks/abc331_a

## 考察
次の日にちを求めたいので、とりあえず $d$ を $+1$ します。
1か月は $D$ 日までなので、もし $d$ が $D$ を超えていたら、 $d$ を $1$ にして、 $m$ を $+1$ します。
さらに、1年間は $M$ 日までなので、もし $m$ が $M$ を超えていたら、 $m$ を $1$ にして、 $y$ を $+1$ します。


## コード
```Python
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$ 個以上買うための最小の金額を求めてください。

https://atcoder.jp/contests/abc331/tasks/abc331_b

## 考察
卵のパックは、どれも $N+1$ 個以上買う必要はありません。(本当はもっとこの上限の値は下げられますが、今回はそこまで計算しなくてもACできるのでこれでいきます)
制約を見てみると $N \leq 100$ なので、 卵のパックの数をそれぞれ $0$ ～ $N$ 個の中で全探索しても探索回数は $100^3=10^6$ 回程度です。

## コード
```Python
# 全体の値段
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$ よりも大きな要素すべての和を求めてください。

https://atcoder.jp/contests/abc331/tasks/abc331_c

## 考察
愚直な解き方としては、各 $i=0,1,\cdots N$ について、 $A$ の中で $A_i$ よりも大きな数の和を出力する下のようなコードがあります。
```Python
"""愚直な解き方"""
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してしまいます。
なのでどうにか工夫したいです。

方針として、下のような辞書をつくりたいです。
```Python
"""辞書をつくりたい"""
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は、普通の辞書と同じように扱えて、なおかつ常にキーが小さい順にソートされるような辞書です。
```Python
"""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することができます。


## コード
```Python
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)$ の長方形の中にある黒マスの個数はいくつですか？

https://atcoder.jp/contests/abc331/tasks/abc331_d

## 考察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)$ のときの図です。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/b04b74e3-5420-29f2-b3bf-d963ee91ba1c.png)

まず、$f(4,5)$ で、 $(0,0)$ から $(3,4)$ までの長方形の中にある黒マスの数を求めます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/e7781a07-d1e4-a965-9a24-2850bcea85ed.png)

無駄に数えてしまっている部分があるので、その部分を引き算したいです。
$f(4,2)$ と $f(1,5)$ を求めて引きます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/b6d52cb1-df01-898d-62a9-a993f2246f5d.png)

$f(1,2)$ の部分を2度引いてしまっているので、帳尻合わせで足しておきます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/86abb4f7-8944-847a-a81f-808e80007dc3.png)

まとめると、$(A,B,C,D)=(1,2,3,4)$ のとき、答えは $f(4,5)-f(4,2)-f(1,5)+f(1,2)$ です。
ここまでの解答コードはこんな感じになります。

```Python
"""考察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$ のときの答えは先に計算してメモしておきましょう。
下のコードでは、いもす法を縦と横の両方にすることで求めています。

```Python
# 半開区間[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$ とします。


![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/921f76d8-2394-6ee4-4300-8f612d6ffa18.png)


オレンジの部分は、$N \times N$ の$P$ 内にある黒マスの数を sum とすると、 $sum \times q_i\times q_j$ です。この変数は、考察2で求めています。

緑・灰色の部分は、次のコードで求められます。
```Python
green = P_acc[mi][N] * qj
gray = P_acc[N][mj] * qi
```

青色の部分は P_acc をそのまま利用して求められます。

まとめると、解答コードは次のようになります。

## コード
```Python
# 半開区間[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

## 問題
主菜と副菜の組み合わせのなかで、最も高い金額になるものはいくらですか？
ただし、食べ合わせの悪い組み合わせを選ぶのはナシです。

https://atcoder.jp/contests/abc331/tasks/abc331_e

## 考察
副菜のリスト $B$ を大きい順にソートします。
主菜のリスト $A$ の要素 $a$ それぞれについて、ソートされたリスト $B$ の1つめの要素 $B_0$ と組み合わせて定食をつくります。
もし食べ合わせが悪いなら、$B_1$ と組み合わせ、食べ合わせが良くなるまでこれを繰り返します。
食べ合わせが良いとき、変数ansを更新します。
この $a$ について、さらに副菜リスト $B$ の要素を見る必要はありません。 $B$ はソートされていて、この先はより安い金額しか現れないからです。
これを $A$ の要素すべてについて行えばokです。
一見すると計算量が $O(NM+MlogM)$ に見えますが、食べ合わせの悪い組み合わせは $L$ 個なので、全体で $O(N+MlogM+L)$ です。

## コード
```Python
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)
```