見出し画像

【ABC356】緑コーダーの競プロ参加記録 #21

今回はAtCoder Beginner Contest 356。使用言語はPythonです。


A問題 - Subsegment Reverse

コンテスト中の提出
https://atcoder.jp/contests/abc356/submissions/54078759


  • 正整数$${N, L, R}$$が与えられる。

  • 数列$${A = (1, 2, …, N)}$$。

  • $${A}$$の$${L}$$項目から$${R}$$項目を反転した後の$${A}$$を出力してね。


元の数列$${A}$$を用意して、スライスを使って反転すると比較的シンプルな実装で解くことができます。

スライスで反転するときは$${[: :-1]}$$と書きます。

""" AC コード """
N, L ,R = map(int, input().split())
A = [i + 1 for i in range(N)]
A[L - 1 : R] = A[L - 1 : R][::-1]
print(*A)

数列$${A}$$を用意せず、print関数の end キーワードを利用して解くこともできます。

""" ACコード """
N, L ,R = map(int, input().split())
for i in range(1, L):
    print(i, end=" ")
for i in reversed(range(L, R + 1)):
    print(i, end=" ")
for i in range(R + 1, N + 1):
    print(i, end=" ")

上のようにrange関数を3回使って、操作後の$${A}$$を作ってもよいです。Python のリストは '+' 演算子で連結することができます。

""" ACコード """
N, L ,R = map(int, input().split())
A = [i for i in range(1, L)]
A += [i for i in reversed(range(L, R + 1))]
A += [i for i in range(R + 1, N + 1)]
print(*A)


B問題 - Nutrients

コンテスト中の提出
https://atcoder.jp/contests/abc356/submissions/54082709


  • $${N}$$品の食品と、$${M}$$種類の栄養素がある。

  • $${i}$$番目の食品では、栄養素$${j}$$が$${X_{i,j}}$$とれる。

  • $${j}$$番目の栄養素は$${A_j}$$以上ほしい。

  • すべての栄養素について、ちゃんと栄養がとれているか判定してね。


二次元リスト$${X}$$を全探索して、栄養素ごとに摂取量の総和を求めます。

その後、必要量$${A}$$と比較し、必要量を満たしていない栄養素が1つでもあれば答えは 'No' です。

""" 入力を受け取る """
N, M = map(int, input().split())
A = list(map(int, input().split()))
X = [list(map(int, input().split())) for _ in range(N)]

""" 栄養素ごとに、摂取量を調べる """
res = [0] * M
for i in range(N):
    for j in range(M):
        res[j] += X[i][j]

""" 栄養素ごとに、必要量を満たしているか調べる """
for i in range(M):
    if A[i] > res[i]:
        print('No')
        exit()
print('Yes')


C問題 - Keys

コンテスト中の提出
https://atcoder.jp/contests/abc356/submissions/54100941


  • $${N}$$本の鍵がある。正しい鍵とダミー鍵が混ざっている。

  • $${K}$$本の正しい鍵を挿すと開くドアに対して、$${M}$$回のテストをした。

  • テスト: $${C_i}$$本の鍵$${A_{i,0},A_{i,1},…,A_{i,C_i}}$$を挿して、結果は$${R_i}$$。

  • 正しい鍵とダミー鍵の組み合わせとしてテストに矛盾しないものが何通りあるか出力してね。


「各鍵が正しいかダミーかの組み合わせは$${2^N}$$通り考えられます」と、問題文に書いてあります。

$${1 \le N \le 15}$$であり$${2^N}$$通り試すことができそうなので、bit全探索をします。

ビットマスクの右から$${i}$$桁目が$${1}$$であればカギ$${i}$$は正しい鍵、$${i}$$桁目が$${0}$$であればカギ$${i}$$はダミー鍵であるとします。

ビットマスクは 0-index として扱いたいので、入力を受け取る際に$${A}$$の要素を$${-1}$$しておきます。

""" 入力を受け取る """
N, M, K = map(int, input().split())
C, A, R = [],[],[]
for _ in range(M):
    c, *a, r = input().split()
    C.append(int(c))
    A.append([int(i) - 1 for i in a])    # 0-index に変換
    R.append(r)

さて、「テストに矛盾しない」とはどういうことなのでしょうか。
これは、以下の2通りが考えられます。

  • 正しい鍵を$${K}$$本以上挿していて、テスト結果が 'o' (開く)

  • 正しい鍵を$${K}$$本未満しか挿しておらず、テスト結果が 'x' (開かない)

これをすべてのテストについて判定して、問題が無ければ答えを$${+1}$$します。

ビットマスク (mask) の右から$${i}$$桁目が$${1}$$であるかどうかを判定するためには、「ビットマスクを$${i}$$回右シフトした値と$${1}$$を$${\mathrm{AND}}$$演算した結果」が$${1}$$であるかを確認します。

""" 鍵の組み合わせを bit全探索 """
ans = 0
for mask in range(1 << N):
    ok = True
    for i in range(M):
        cnt = 0
        for a in A[i]:
            if (mask >> a) & 1 == 1:
                cnt += 1
        if R[i] == 'o' and cnt < K:
            ok = False
        if R[i] == 'x' and cnt >= K:
            ok = False
    if ok:
        ans += 1
print(ans)

なお、ビットマスクの$${1}$$である桁が$${K}$$個より少なければ正しい組み合わせでないように思えますが、テスト結果が全て 'x' であれば矛盾しないため注意が必要です。

CPython で提出すると TLE になりました。
PyPyで提出をすると 400 ms 程度で AC になります。


問題文を読み解くのが難しい問題だったように思います。
何が分かっていればよいのか理解するまで、私は10分くらいかかりました。


D問題 - Masked Popcount

コンテスト中の提出
https://atcoder.jp/contests/abc356/submissions/54121087


  • 整数$${N, M}$$が与えられる。

  • $${\displaystyle\sum_{k =0}^N \mathrm{popcount}(k  \&  M)}$$を出力してね。$${\pmod {998244353}}$$


ビット演算がメインテーマの問題では、桁ごとに独立して考えるケースが多いです。

「$${0}$$以上$${N}$$以下の整数」のうち「桁についての何かしらの条件」を満たす「整数の個数」を効率よく数え上げるアルゴリズムとして、桁DPがあります。

$${\mathrm{popcount}}$$はまさに「桁についての何かしら」の情報です。

今回の問題では、「$${0}$$以上$${N}$$以下の整数」のうち「$${ \mathrm{popcount}(k  \& M) = j}$$」である「整数$${k}$$の個数$${dp[j]}$$」が分かれば、すべての$${j}$$について$${dp[j] \times j}$$の総和が答えになります。


桁DP

桁DP では、左の桁から1つずつ数を確定させていきます。

その際、現在見ている桁の段階で「$${N}$$以下であることが確定しているもの」と、「$${N}$$以下であることが確定していないもの」を分けて考える必要があります。

つまり、以下のようなDPを考えます。

$$
dp[i][j][smaller]: 左からi桁目の時点で、\mathrm{popcount}(k  \& M) = j である、smallerな整数kの個数 \\
smaller = \begin{cases}
0 : N以下であることが未確定 \\
1 : N以下であることが確定 \\
\end{cases}
$$

0 桁目時点では$${N}$$以下であることが確定していません。よって$${dp[0][0][0] = 1}$$です。

$${i}$$桁目時点で$${N}$$以下であることが確定している、すなわち$${smaller = 1}$$であるとき$${i + 1}$$桁目の数字はなんでもOKです。
(二進数なので選べる数字は$${\{0,1\}}$$のどちらかです。)

$${i}$$桁目時点で$${N}$$以下であることが確定していない、すなわち$${smaller = 0}$$であるとき、

  • $${i + 1}$$桁目に「$${N}$$の$${i + 1}$$桁目と同じ数字」を選ぶ場合は、まだ$${N}$$以下であることが確定しないままです。$${(smaller = 0)}$$

  • $${i + 1}$$桁目に「$${N}$$の$${i + 1}$$桁目より小さい数字」を選ぶ場合は、$${N}$$以下であることが確定します。$${(smaller = 1)}$$

  • $${N}$$の$${i + 1}$$桁目より大きい数字は選べません

さらに今回の問題では、$${i+1}$$桁目に$${1}$$を選択した場合に、$${M}$$の$${i+1}$$桁目が$${1}$$であれば$${\mathrm{popcount}}$$が$${+1}$$します。

""" ACコード """
N, M = map(int, input().split())
mod = 998244353
"""
dp[i][j][smaller]:  左からi桁目の時点で、
                    popcount(k & M) = j である、
                    smaller(N 以下)な整数 k の個数
"""
D = 60
dp = [[[0] * 2 for _ in range(D + 2)] for _ in range(D + 2)]
dp[0][0][0] = 1
for i in range(D):
    d = D - (i + 1)
    for j in range(D):
        if (N >> d) & 1 == 1:
            """ i + 1 桁目に 0 を選んだ場合 (N 以下が確定) """
            dp[i + 1][j][1] += dp[i][j][1] + dp[i][j][0]

            """ i + 1 桁目に 1 を選んだ場合 """
            if (M >> d) & 1 == 1:
                dp[i + 1][j + 1][0] += dp[i][j][0]
                dp[i + 1][j + 1][1] += dp[i][j][1]
            else:
                dp[i + 1][j][0] += dp[i][j][0]
                dp[i + 1][j][1] += dp[i][j][1]
        else:
            """ i + 1 桁目に 0 を選んだ場合 """
            dp[i + 1][j][0] += dp[i][j][0]
            dp[i + 1][j][1] += dp[i][j][1]

            """ i + 1 桁目に 1 を選んだ場合 """
            if (M >> d) & 1 == 1:
                dp[i + 1][j + 1][1] += dp[i][j][1]
            else:
                dp[i + 1][j][1] += dp[i][j][1]
        dp[i + 1][j][0] %= mod
        dp[i + 1][j][1] %= mod
ans = 0
for j in range(D + 1):
    ans += (dp[D][j][0] + dp[D][j][1]) * j
    ans %= mod
print(ans)

今回は$${i}$$が1つ前のものしか参照しないためDPのリストを分割できます。
さらに、$${\mathrm{smaller}}$$を別々のリストに分けることで、一次元のリストで解くこともできます。

$$
dp[j]: i桁目の時点で、\mathrm{popcount}(k  \& M) = j である整数kの個数 \\

dp0, ndp0 : N以下であることが未確定 \\
dp1, ndp1 : N以下であることが確定 \\
$$

N, M = map(int, input().split())
mod = 998244353
K = N.bit_length()
""" 
dp[j]: popcount(k & M) = j である整数 k の個数
(dp0, ndp0: N 以下であることが未確定) 
(dp1, ndp1: N 以下であることが確定)
"""
dp0 = [0] * (K + 2)
dp1 = [0] * (K + 2)
dp0[0] = 1
for i in reversed(range(K)):
    ndp0 = [0] * (K + 2)
    ndp1 = [0] * (K + 2)
    for j in range(K):
        if (N >> i) & 1 == 1:
            """ 現在の桁に 0 を選んだ場合 (N 以下が確定) """
            ndp1[j] += dp1[j] + dp0[j]

            """ 現在の桁に 1 を選んだ場合 """
            if (M >> i) & 1 == 1:
                ndp0[j + 1] += dp0[j]
                ndp1[j + 1] += dp1[j]
            else:
                ndp0[j] += dp0[j]
                ndp1[j] += dp1[j]
        else:
            """ 現在の桁に 0 を選んだ場合 """
            ndp0[j] += dp0[j]
            ndp1[j] += dp1[j]

            """ 現在の桁に 1 を選んだ場合 """
            if (M >> i) & 1 == 1:
                ndp1[j + 1] += dp1[j]
            else:
                ndp1[j] += dp1[j]
        ndp1[j] %= mod
        ndp1[j + 1] %= mod
    dp0 = ndp0
    dp1 = ndp1
ans = 0
for i in range(K + 1):
    ans += (dp0[i] + dp1[i]) * i
    ans %= mod
print(ans)

個人的には、3次元DPよりこちらの方が添え字のミスが少ないうえに、目が滑らないので分かりやすい気がします。

公式解説などを見るに、桁DPが想定解ではない(オーバーすぎる?)解法であるように思えるので、分かりやすいと思うものを選んでください。


E問題 - Max/Min

upsolve
https://atcoder.jp/contests/abc356/submissions/54152623


  • 長さ$${N}$$の数列$${A}$$が与えられる。

  • $${\displaystyle\sum_{i=1}^{N-1}\displaystyle\sum_{j=i+1}^{N}\bigg\lfloor \frac{\max(A_i, A_j)}{\min(A_i, A_j)}\bigg\rfloor}$$を出力してね。


ある$${A_i}$$について、$${{\Large\lfloor \frac{A_j}{A_i}\rfloor} = x}$$となる$${A_j}$$の範囲は以下の通りです。

$$
xA_i \le A_j < (x+1)A_i
$$

よって、各$${A_i}$$について$${x}$$を全探索し、$${xA_i \le A_j < (x+1)A_i}$$を満たす$${A_j}$$の個数$${n}$$が分かれば、答えに$${n \times x}$$を加算します。

「$${xA_i \le A_j < (x+1)A_i}$$を満たす$${A_j}$$の個数」は、整数$${k}$$の個数を$${Cnt[k]}$$とするリストの累積和を前計算しておくと$${O(1)}$$で取得できます。

""" 累積和をとる """
from itertools import accumulate
N = int(input())
A = list(map(int, input().split()))
M = 10 ** 6
Cnt = [0] * (M + 1)
for a in A:
    Cnt[a] += 1
acc = list(accumulate(Cnt))

$${A_i}$$の重複分はまとめて計算したほうが分かりやすいです。

$${{\Large\lfloor \frac{A_j}{A_i} \rfloor} = 1}$$すなわち$${x = 1}$$であるとき、$${A_i = A_j}$$であるケースと$${A_i \not = A_j}$$であるケースを分けて計算する必要があることに注意が必要です。

$${A_i = A_j}$$であるケースは$${_{Cnt[A_i]}C_2}$$通りです。

""" 同じ数はまとめて計算する """
ans = 0
for a in set(A):
    X = M // a
    for x in range(1, X + 1):
        l = a * x
        r = min(M, a * (x + 1) - 1)
        n = acc[r] - acc[l - 1]
        if x == 1:
            ans += Cnt[a] * (Cnt[a] - 1) // 2
            n -= Cnt[a]
        ans +=  n * x * Cnt[a]
print(ans)

$${A_i}$$が大きくなるごとに、全探索する$${x}$$の範囲が小さくなっていくので、トータルの計算量は$${O(M \log{N})}$$になるらしいです。

調和級数がなんやかんやです。(公式解説


あとがき

今回は ABCD 4完 1ペナ 80分で 2892位、1087 perf でした。

適正レートだなぁ。

ではまた~。


【更新履歴】
2024/06/02 (18時頃):記事投稿。


いいなと思ったら応援しよう!