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

AtCoder Beginners Contest 318 (ABC318) をPythonで解きました。
見やすいコードを書く練習も兼ねてます。

A - Full Moon

問題ページ

考察

変数 $ans$ を用意して、最初は $ans=0$ とします。
$M$ 日目、 $M+P$ 日目、 $M+2P$ 日目、$\cdots$ と見ていきます。
もし今見ている日にちが $N$ 以下のときは、変数 $ans$ を $1$ だけ増やして、次の日にちを見に行きます。
もし今見ている日にちが $N$ を超えていたら、日にちを $P$ ずつ増やしていた操作をbreak文などで打ち切って、変数 $ans$ を出力します。
下のコードはwhile文をつかっています。別解にもその他の書き方があるので、参考にしてみてください。

コード

N, M, P = map(int, input().split())

ans = 0
while M <= N:
    ans += 1
    M += P
print(ans)

別解

for文をつかった解法

for文をつかって、 $M$ 日目、 $M+P$ 日目、　$M+2P$ 日目、 $\cdots$ と日付を見ていきます。 rangeの引数には、どうやっても見ることのないくらい大きな数字を入れておき、見ている日にちが $N$ を超えた時点でbreak文をつかってこのループ処理を打ち切ります。

# for文をつかった解法
N, M, P = map(int, input().split())
ans = 0
for i in range(5000000):
    day = M + P * i
    if day <= N:
        ans += 1
    else:
        break
print(ans)

range関数の第3引数をつかった解法

range関数は、たとえば「 $3$ から $14$ までの整数を $3$ 個とばしで見る」みたいなことができます。

指定するのは $3$ つで、「スタートの数」「おわりの数」「何個とばしでみるか」です。「おわりの数」は、本来の数から $1$ だけ足しておかないといけません。

つかいかたはこんな感じ。

# 3から14までの整数を3個とばしで見る。
for i in range(3, 15, 3):
    print(i)  # 3, 6, 9, 12 が順に出力される

# 5から22までの整数を4個とばしで見る。
for i in range(5, 23, 4):
    print(i)  # 5, 9, 13, 17, 21 が順に出力される

この性質を用いて、今回の問題を解くことができます。

# range関数の第3引数をつかった解法
N, M, P = map(int, input().split())
ans = 0
for day in range(M, N + 1, P):
    ans += 1
print(ans)

len関数をつかった解法(発展)

上の「range関数の第3引数をつかった解法」の発展バージョンです。

今回の問題では、何回ループするかの情報がほしいだけです。これは実際にfor文をつかってループさせなくても、len関数をつかって取得できます。

# len関数をつかった解法
N, M, P = map(int, input().split())
ans = len(range(M, N + 1, P))
print(ans)

こっちで計算する解法(発展)

for文やwhile文をつかってPythonにループ処理を頑張ってもらってもいいですが、自分で計算することもできます。 スタートの日にちが入力で与えられる $M$ によって毎回変わってしまうので、最初の日にちを $0$ 日目で統一するために、次のような変更をします。

• 満月を見られる最初の日は $M$ 日目 → 満月を見られる最初の日は $0$ 日目
• $1$ 日目から $N$ 日目の中で、満月を見られる日の数を求める → $0$ 日目から $N-M$ 日目の中で、満月を見られる日の数を求める

本来は今日がスタートだったのを $M$ 日後の今日に置き換えて、 1-indexedを0-indexedになおしました。
たとえば、「満月は $5$ 日ごとに現れる( $P=5$ )とき、 $0$ 日目から $28$ 日目の中で、満月を見られる日の数を求めてください。ただし、 $0$ 日目に満月が見られるとします。」という問題のとき、満月が見られるのは $0$ 日目、 $5$ 日目、 $10$ 日目 、 $15$ 日目、 $20$ 日目、 $25$ 日目の $6$ 回です。これは、$28//5 + 1=6$ として求めることもできます(//は小数点以下を切り捨てる割り算)。

これをそのままコードにすることで解くことができます。
$N<M$ のときは答えが $0$ になるので注意しましょう。

# こっちで計算する解法
N, M, P = map(int, input().split())
ans = (N - M) // P + 1
if ans < 0:
    ans = 0
print(ans)

B - Overlapping sheets

問題ページ

考察

$101 \times 101$ のマスを用意して、シートで覆った部分はTrue、そうでない部分はFalseということにします。
$N$ 枚のシートが与えられるので、それぞれのシートについて、覆った部分をTrueにしていきます。
$101 \times 101$ のマスの中にあるTrueの数がそのまま答えになります。

コード

N = int(input())
Grid = [[False] * 101 for _ in range(101)]
for _ in range(N):
    a, b, c, d = map(int, input().split())
    for i in range(a, b):
        for j in range(c, d):
            Grid[i][j] = True
ans = 0
for i in range(101):
    for j in range(101):
        if Grid[i][j]: ans += 1
print(ans)

C - Blue Spring

問題ページ

考察

1日周遊パスはどのタイミングでつかってもいいです。
また気持ち的には、通常料金の高い日に1日周遊パスをつかって、通常料金の安い日は1日周遊パスはつかわないようにしたいです。

なので、リスト $F$ で与えられる各日にちの通常料金のリストを、大きい順(降順)にソートします。サンプル1だと、$F=[7,1,6,3,6]$ だったのが、 $F=[7,6,6,3,1]$ になります。

最初の $D$ 日間を見てみます。
もし最初の $D$ 日間は1日周遊パスをつかった方がいいなら、$D$ 枚セットの1日周遊パスを買って、最初の $D$ 日間につかうことにします。
サンプル1だと、 $F=[7,6,6,3,1]$ で、1日周遊パスは $2$ 枚セットで $10$ 円です。最初の $2$ 日間は何もしないと $13$ 円かかってしまうので、$2$ 枚セットの1日周遊パスを $10$ 円で購入します。
これを次の $D$ 日間、さらに次の $D$ 日間、 $\cdots$ と繰り返していき、実際にかかるお金がいくらかを計算すればokです。

下のコードでは、 $D$ 枚セットの1日周遊パスを買ってつかうとき、 $D$ 日間のうち最初の日だけ $P$ 円(=1日周遊パスの値段)、それ以外の日は $0$ 円としています。

コード

N, D, P = map(int, input().split())
F = list(map(int, input().split()))
F.sort(reverse=True)

for left in range(0, N, D):
    right = min(left + D, N)
    cost = sum(F[left:right])
    if cost > P:
        for i in range(left, right):
            F[i] = 0
        F[left] = P
    else:
        break
ans = sum(F)
print(ans)

D - General Weighted Max Matching

問題ページ

考察

考えうる辺の結び方は、 $15 \times 13 \times 11 \times 9 \times 7 \times 5 \times 3 = 2027025$　通りで、これなら全探索できそうです。
DFS全探索で解きます。

再帰関数をつかうので、次の4行を書いておくようにしましょう。

import sys
import pypyjit

sys.setrecursionlimit(10**7)
pypyjit.set_param('max_unroll_recursion=-1')

$N$ が奇数のとき、浮いちゃう頂点が出てきてややこしいので、頂点を1つ増やしてその頂点から出ている辺はすべて重みを $0$ とします。

• $lst[i]$ : 辺の番号が2つ格納されたリスト

となるリスト $lst$ を全パターンつくります。
以下、どのようにつくるかを書いていきます。

格納する辺の番号を $v$ 、つくってる途中のリストを $lst$ とします。このことを $dfs(lst,v)$
と表記します。
$lst$ を左から順に見ていきます。
要素数が $2$ のとき、その2頂点をつなぐ予定なので、そこには $v$ は入れません。
要素数が $0$ のとき、そこに $v$ を入れた新たなリスト $nexlst$ を用意し、$dfs(nexlst, v+1)$ を考えることにします。 ただし、もう要素数が $0$ のところに $v$ は入れません。
要素数が $1$ のとき、そこに $v$ を入れた新たなリスト $nexlst$ を用意し、$dfs(nexlst, v+1)$ を考えることにします。

これで全パターンの $lst$ を列挙することができます。

コード

import sys
import pypyjit

inputs = sys.stdin.readline
sys.setrecursionlimit(10 ** 7)
pypyjit.set_param('max_unroll_recursion=-1')

N = int(inputs())
D = []
for i in range(N - 1):
    row = list(map(int, inputs().split()))
    D.append([0] * (i + 1) + row)
D.append([0] * N)
for i in range(N):
    for j in range(N):
        if i != j and D[i][j] == 0:
            D[i][j] = D[j][i]

if N % 2 == 1:
    for i in range(N):
        D[i].append(0)
    D.append([0] * (N + 1))
    N += 1

ans = 0


def dfs(lst, v):
    if v == N:
        global ans
        now_ans = 0
        for u, v in lst:
            now_ans += D[u][v]
        ans = max(ans, now_ans)
        return

    append_to_zero = False  # 過去に要素数0のところにvを追加したことがあればTrue
    for i in range(len(lst)):
        match len(lst[i]):
            case 2:
                continue
            case 0:
                if append_to_zero:
                    break
                append_to_zero = True
        nex_lst = [t[:] for t in lst]

        nex_lst[i].append(v)
        dfs(nex_lst, v + 1)


dfs([[] for _ in range(N // 2)], 0)
print(ans)

別解

bitDPをつかった解法

$0$ ～ $2^N - 1$ の合計 $2^N$ 個の数を用意します。これらの数は、「2進数で表したときに下から $i$ 桁目が $1$ のとき、頂点 $i$ はすでに他の頂点と結ばれている」という状態を表すことにします。

• $dp[S]$ : グラフの状態が $S$ のときの、辺の重みの総和の最大値

となるdpを考えます。
最初は $dp[0] = 0$ です。

以下、下のコードで読み解くのがしんどくなりそうなところを解説します。

グラフの状態を表す整数 $S$ について、頂点 $i$ がすでに他の頂点と結ばれているかどうかは、次のように判定できます。

if (S >> i) & 1 == 1:
    print('頂点iは他の頂点と結ばれている')
elif (S >> i) & 1 == 0:
    print('頂点iはどの頂点とも結ばれていない')

また、もとのグラフの状態が $S$ で、そこに頂点 $a,b$ を結ぶ辺を追加するのは、次のように書けます。

S |= 1 << a
S |= 1 << b

itertools.combinations()をつかうと、あるリストから組み合わせのパターンをすべて列挙することができます。

from itertools import combinations

lst = [5, 12, 3, -92]
for a, b in combinations(lst, 2):
    print(a, b)

# 出力結果
# 5 12
# 5 3
# 5 -92
# 12 3
# 12 -92
# 3 -92

# bitDPをつかった解法
from collections import deque
from itertools import combinations

N = int(input())
D = []
for i in range(N - 1):
    row = list(map(int, input().split()))
    D.append([0] * (i + 1) + row)
D.append([0] * N)
for i in range(N):
    for j in range(N):
        if i != j and D[i][j] == 0:
            D[i][j] = D[j][i]

dp = [0] * (1 << N)
que = deque()
que.append(0)
while que:
    state = que.popleft()
    P = [i for i in range(N) if (state >> i) & 1 == 0]  # つかえる頂点のリスト
    if len(P) >= 2:
        for a, b in combinations(P, 2):
            nex_state = state
            nex_state |= 1 << a
            nex_state |= 1 << b
            if dp[nex_state] == 0:
                que.append(nex_state)
            dp[nex_state] = max(dp[nex_state], dp[state] + D[a][b])
print(max(dp))

E - Sandwiches

問題ページ

考察 (パート1)

次の問題を考えます。

長さ $N$ の整数列 $A=(A_1,A_2, \cdots ,A_N)$ が与えられます。以下の条件をすべて満たす正整数組 $(i,j,k)$ の個数を求めてください。

• $1 \leq i < j < k \leq N$
• $A_i = A_k$

もとの問題文から、$A_j \neq A_k$ の条件を取り除いただけです。この問題をいったん解いてみます。

サンプル3の $A=(9, 7, 11, 7, 3, 8, 1, 13, 11, 11, 11, 6, 13)$ を考えてみます。
たとえば $k=10$ (右から $4$ 番目)のときを考えます。
$A_i=A_{10}$ となる $i$ は、 $i=3,9$ です。
$(i,k)=(3,10)$ のとき、$j=4,5,6,7,8,9$ で $6$ 通りです。これは、 $k-i-1=6$ で算出できます。
$(i,k)=(9,10)$ のときは、$0$ 通りです。

次に $k=11$ のときを考えます。
$A_i=A_{11}$ となる $i$ は、 $i=3,9,10$ です。
$(i,k)=(3,11)$ のとき、 $k-i-1=7$ で $7$ 通りです。
$(i,k)=(9,11)$ のとき、 $k-i-1=1$ で $1$ 通りです。
$(i,k)=(10,11)$ のときは、$0$ 通りです。

今回はもう値が $11$ の場所がこの右側にないですが、もし右側にまだまだ $11$ の値があったら、そのたびに $3$ や $9$、 $10$ を $k$ から引くことになってしまいます。
毎回計算しなくても、 $k \times (左側で11の出た回数) - (11が出てきたインデックスの総和)$ で求められます。
この計算をするときに必要なのは、今までに左側で何回 $11$ が出てきたか、$11$ が出てきたインデックスの総和、の2つだけです。
なので、次の2つのリストを用意します。

• $cnt[i]$ : 数字 $i$ が登場した回数
• $value[i]$ : 数字 $i$ が登場したインデックスの総和

この2つのリストを管理しながら、$k=1$ から $k=N$ まで計算していきます。
これで、最初につくった問題を解くことができました。

考察 (パート2)

次の問題を考えます。

長さ $N$ の整数列 $A=(A_1,A_2, \cdots ,A_N)$ が与えられます。以下の条件をすべて満たす正整数組 $(i,j,k)$ の個数を求めてください。

• $1 \leq i < j < k \leq N$
• $A_i = A_k$
• $A_j = A_k$

もとの問題文の$A_j \neq A_k$ の条件を $A_j=A_k$ に変えました。
パート1での問題の答えから、この問題の答えを引くことで、もとの問題の答えになります(補集合の考え方です)。

$A_i$ と同じ値の数を $c$ とすると、その数について $A_i=A_j=A_k$ になるパターンの数は、$_cC_3=c(c-1)(c-2)/6$ です。

これで、もとの問題を解くことができました。

コード

N = int(input())
A = list(map(lambda a: int(a) - 1, input().split()))

ans = 0
cnt = [0] * N  
value = [0] * N
for i, a in enumerate(A):
    c = cnt[a]
    v = value[a]
    if c >= 1:
        ans += i * c - v - c
    cnt[a] += 1
    value[a] += i

for c in cnt:
    if c <= 2:
        continue
    ans -= c * (c - 1) * (c - 2) // 6

print(ans)