2022-09-13

ABC127E

AtCoder Python 組み合わせマンハッタン距離駒の配置

問題

atcoder.jp $N \times M$ マスに配置する $K$ 個の駒の間のマンハッタン距離の総和を，とりうる配置のパターンでさらに総和する問題です。

解説

公式解説を読んでも理解に苦しめられたので，少しだけ噛み砕いて記事にします。
そしてなんと，今回からMarkdown記法で書くことにしました。ほとんど更新してなくてすみません。

2点 $P(x_i, y_i), Q(x_j, y_j)$ 間のマンハッタン距離を場合の数だけ掛け合わせると良いです。これについて，具体例をもとに３つのSTEPに分けて理解していきたいと思います。

具体例 $(N, M, K) = (2, 2, 3)$ で考える

$2 \times 2$ のマスに3個の点を配置する場合を考えます。

STEP 1

点の配置の仕方は図のようになります。点を区別するために番号を振っています。
各配置におけるマンハッタン距離の総和を４通りの場合で総和をとると答えになります。

STEP 2

すべての配置を列挙するのは難しいので，２点間の距離のみに着目して，それが答えに何回足されるかを考えます。残りの点の配置の仕方は，残りの2マスの中から1マスを選ぶので ${}_{2} C_{1}$ 通りです。一般化すると， ${}_{NM-2} C_{K-2}$ 通りになります。 $N \times M \leqq 2 \times 10^{5}$ の制約では， $O(NM)$ で算出しても間に合います。
PQの選び方を列挙したら，それぞれのマンハッタン距離を算出し，場合の数をかけて合算します。

ここで， $N \times M$ のグリッドから2点を選ぶ場合の数は ${}_{NM} C_{2}$ であり，これを全探索してしまうと計算量は $O(N^{2}M^{2})$ になります。もし，pythonで書くならば，

for i in range(N*M-1):
    for j in range(i+1, N*M):
        ...

のようになると思います。これでは間に合いません。そこで次の考え方に移ります。

STEP 3

マンハッタン距離の計算において， $x$ , $y$ は独立しているので，今回は分離して考えます。 x座標の差を $dx$ ，y座標の差を $dy$ で表しています。
$dx=0$ のときは， $dx$ に何をかけても0なので，考慮する必要はありません。
$dx=i$ の選び方とき，2点PQの選び方は次の3つの自由度から定められるので， $(N-i) M^{2}$ 通りになります。
この計算は，ループ処理で1からN-1まで距離を求めればよいので，計算量 $O(N)$ で求まります。

1つ目の点の $x$ 座標を 1 から $N-x$ の中からどれにするか
行の選び方（クリックすると展開されます）
ここの"行の選び方"は， $1 \leqq x_i \leqq x_j \leqq N$ のうち， $x_j - x_i = dx$ となる $(x_i,\ x_j)$ の選び方です。
例えば， $x=(1, 2, 3)$ のうち，差が 1 となる選び方は， $(x_i,\ x_j) = (1, 2), (2, 3)$ の2通りとなります。
1つ目の点の $y$ 座標を 1から $M$ の中からどれにするか
2つ目の点の $y$ 座標を 1から $M$ の中からどれにするか

y座標についても，x座標と同様に計算ができます。最後に，STEP 2 で求めた残りの点の場合の数を考慮して， $x$ の寄与と $y$ の寄与の総和を求めると答えになります。

$\displaystyle ans = \biggl( \sum_{i=1}^{N-1} {i(N-i)M^{2}} + \sum_{i=1}^{M-1} {i(M-i)N^{2}} \biggr) \times {}_{NM-2} C_{K-2}$

コード

提出コード: リンク
MODnCr: GitHub

def main():
    MOD = 1_000_000_007
    N, M, K = map(int, input().split())
    
    ans= 0
    # axis: x
    for d in range(1, N):
        ans += d * (N-d)*M**2
        ans %= MOD
    # axis: y
    for d in range(1, M):
        ans += d * (M-d)*N**2
        ans %= MOD
 
    nCr = MODnCr(mod=MOD)
    ans *= nCr(N*M-2, K-2)
    ans %= MOD
    print(ans)
    return

ABC256F

ABC256F の解説

atcoder.jp

解答コード

https://atcoder.jp/contests/abc256/submissions/32580103 (1534 ms)

問題

タイトル通り，数列Aの累積和Bの累積和Cの累積和Dを求める問題です．

解説

概要図です．

概要図の内容は，公式解説（by kyopro_friends）を図解したものになります．
Dx = ... の式の導出は，公式解説を参照してください m(_ _)m．

この問題のポイントです．
Aの点更新は，Bの区間更新になり，
Bの点更新は，Cの区間更新になり，
Cの点更新は，Dの区間更新を意味するので，
愚直な実装では到底間に合いません．

そのため，A→B→C→Dが愚直に計算できないので，A→Dを直接計算できるような，第三の配列E, F, Gを用意することを考えます．

ただし，E, F, Gの第x項をクエリの度に累積和をとって，計算量O(N)で算出するようでは間に合わないので，区間和をO(logN)で算出できるBITを用います．

これは，Aの点更新にもO(logN)で計算できるので間に合います．
点更新をする際は，「a を v に置き換える」計算は，「a にv-a を足す」こと再現しています．

コードを提出リンクより参照できますが，あらゆるケースでBITをテストしているわけではなく，動作に確証がないので，もっと整備されたコードを各自参考にしてください．
m(_ _)m．

Parameters:

------------

bit0 : BITのクラス

概要図の配列Exを計算するためのBIT．x=xのときのExの値は，変数e0．

bit1 : BITのクラス

概要図の配列Fxを計算するためのBIT．x=xのときのFxの値は，変数e1．

bit2 : BITのクラス

概要図の配列Gxを計算するためのBIT．x=xのときのGxの値は，変数e2．

class BasicBIT():
    """
    点更新のみBIT
    """
    ...(中略)...

MOD = 998244353
N, Q = map(int, input().split())
A = [0] + list(map(int, input().split()))  # 1-based indexing
bit0 = BasicBIT(A.copy())
bit1 = BasicBIT([a*i for i, a in enumerate(A)])
bit2 = BasicBIT([a*i*i for i, a in enumerate(A)])
for _ in range(Q):
    line = input()
    if line[0] == '1':
        _, x, v = map(int, line.split())
        a = bit0[x]
        bit0.add(x, v-a)
        bit1.add(x, x*(v-a))
        bit2.add(x, x*x*(v-a))
    if line[0] == '2':
        _, x = map(int, line.split())
        e0 = bit0.sum(x+1)
        e1 = bit1.sum(x+1)
        e2 = bit2.sum(x+1)
        ans = (e2 - (2*x+3)*e1 + (x+1)*(x+2)*e0)//2 % MOD
        print(ans)

以上，ABC256Fの解法でした．

＜参考解説リンク＞

・公式解説（by kyopro_friends）

2022-05-11

ABC116D

AtCoder Python

ABC116D の解説

atcoder.jp

解答コード

https://atcoder.jp/contests/abc116/submissions/31604534

問題

いくつかの種類に分かれているN個の寿司の中からK個の寿司を食べた時のコスト関数を最大化する問題です．

コスト = おいしさの総和 + (食べた種類)^2

解説

公式解説の想定解は，初めに美味しいK個の寿司を食べると仮定した貪欲法を試みますが，初めにすべての種類の寿司を食べると仮定して実装しました．こうすることで，種類ボーナスの管理が簡単になると思います．

概要図です．

各種類1位にのネタをまず食べると仮定して配列A，
2位以下ネタをまとめて配列Bとします．

その後，Bの中にある寿司Bi（i=1, 2, ...）を食べた方が良いのかを判定していきます．
まだ食べる予定の寿司の個数がKに満たなければ，無条件でBiを食べることができます．

しかし，K個の寿司が食べる予定になっていれば，Aの中で最も不味いものを残ったBの中で最も美味しいものに入れ替えることで，種類ボーナスを上回る満足度が得られるかを貪欲に判定していきます．

あとは書くだけです．
概要図とソートの昇順降順が逆だったりしますが，そのままなので詳細は割愛します．

Parameters:
-----------
x : int
種類の数です．
capa : int
capacity のことです．あといくつ食べる余裕があるか．

# ABC116D

N, K = map(int, input().split())
T = [[] for _ in range(N)]
for _ in range(N):
    t, d = map(int, input().split())
    T[t-1].append(d)
A = []
B = []
for i in range(N-1, -1, -1):
    if T[i]:
        T[i].sort()
        A.append(T[i].pop())
        B.extend(T[i])
A.sort(reverse=True)
B.sort(reverse=True)
B.append(-1)
A = A[:K]
x = len(A)
capa = K - x
for i, b in enumerate(B):
    if capa:
        capa -= 1
    else:
        if A[-1]+x**2 < b+(x-1)**2:
            A.pop()
            x -= 1
        else:
            break
ans = sum(A) + sum(B[:i]) + x**2
print(ans)

以上，ABC116Dの解法でした．

＜参考解説リンク＞

公式: https://img.atcoder.jp/abc116/editorial.pdf

2022-05-09

ABC250C

AtCoder Python

ABC250C の解説

atcoder.jp

解答コード

https://atcoder.jp/contests/abc250/submissions/31527138

問題

数列の数字を指定して右側の数字とswap する問題です．

解説

ゴールデンウィーク最終日にABC250へ参加した方，大変お疲れ様でした．
m(_ _)m

x のインデックスを計算量O(N) で参照すると間に合わないので，ある数字x がどのインデックスにいるのかを参照する配列を用意する必要があります．

概要図です．

計算量O(1)でインデックスを参照するために，配列（上図の窓口A）を用意する必要がありますが，そのせいで一回のswap 操作の処理が若干複雑になります．
その操作の流れを概要図で示しています．
配列をswap すると同時に，swap 後の新たな住所を窓口Aで住所変更する必要があります．

では，コードに移ります．プログラムはゼロインデックスで書いています．
標準入力と二つの窓口を用意します．

N, Q = map(int, input().split())
A = list(range(N))
B = list(range(N))

１回の操作について解説します．
変更したい数字x を読み取ります．

x = int(input()) - 1

x の場所iを窓口Aに問い合わせます．
i の右隣j を定義します（i が右端の場合，j は左隣）．

i = A[x]
j = i+1 if i+1 < N else i-1

swap していませんが，転出届を窓口Aに提出します．
B[i] は変数x そのものなので，xと書いても良いです．上下揃っているほうが見やすいので，B[i]で書いてます．

A[B[i]] = j
A[B[j]] = i

引っ越しをします．

B[i], B[j] = B[j], B[i]

上記の流れをfor 文で繰り返せば，OKです．
最後にゼロインデックスを解消して出力します．

B = [b+1 for b in B]
print(*B)

以上，ABC250Cの解法でした．

＜参考解説リンク＞

2022-04-28

ABC249C

AtCoder Python

ABC249C の解説

atcoder.jp

解答コード

https://atcoder.jp/contests/abc249/submissions/31307760 （シンプル）
https://atcoder.jp/contests/abc249/submissions/31220829 （numpy）
https://atcoder.jp/contests/abc249/submissions/31221873 （DFS，解説しない）

問題

任意の文字列を選び，アルファベットの個数を数える問題です．

解説

考え方について，概要図を投下します．

文字列の選択の仕方についてbit全探索を行います．そのときに，アルファベットの出現回数がKになる個数をカウントし，最後のそのカウントの最大値を求めれば良いです．

bit全探索を行うための処理ですが，こちらの関数で実行しました．
bit全探索の記述方法はたくさんあります．
itertoolsのproductや，itertoolsのcombinationsを使うことも可能です．
長くなるので，ここでの言及は避けます．
docstring内にあるExample をみてもらえれば，処理の概要は理解してもらえると思います．
全て0で始まる初めのイテレーションを避けたいときに変更が容易なので，最近はこのような自作関数（またはクラス）で使用してます．
解説のために少し手直しました \(^o^)/．

def bitrange(n, start=0, stop=None):
    """
    make bit sequences. 

    Parameters
    ----------
    n : int
        bit length

    Yields
    ------
    seq: list

    Example
    -------
    >>> for seq in bitrange(2):
    >>>    print(seq)
    [0, 0]
    [1, 0]
    [0, 1]
    [1, 1]
    """

    def countup():
        for i in range(n):
            seq[i] += 1
            seq[i] %= 2
            if seq[i] == 1:
                break

    seq = [(start>>j)&1 for j in range(n)]
    if stop is None: 
        stop = 1<<n
    for _ in range(start, stop): 
        yield seq
        countup()

標準入力を読み込みます．Sは長さ26の配列で管理しても，文字列を数字に置き換えた状態のどちらで管理しても良いと思います．
ansのとりうる最小値として暫定的に0で定義しておきます．

N, K = map(int, input().split())
S = [[ord(s)-97 for s in input()] for _ in range(N)]
ans = 0

bit探索を実行します．使用する（seqが1のところの）文字列に含まれるアルファベットの出現個数を+1します．
最後に，count関数でK個の配列数を数えて，暫定的な答えと大小比較します．

for seq in bitrange(N):
    bucket = [0]*26
    for bit, Si in zip(seq, S):
        if bit:
            for si in Si:
                bucket[si] += 1
    cnt = bucket.count(K)
    ans = max(ans, cnt)
print(ans)

続いて，numpyの解法について説明します．
numpyでは，Sの文字列を長さ26のベクトルとして扱います．
説明のために3つのベクトルのうち任意の個数選んだベクトルの和を考えたいと思います．
このとき表は次のようにして求められます．

これを今回のN個の長さ26のベクトルついて考えます．
numpyのインポートと長さ26のN個のベクトルを読み込みます．

N, K = map(int, input().split())
S = np.zeros((N, 26), int)
for i in range(N):
    for s in input():
        S[i, ord(s)-97] = 1

初期値は全て0で定義して，表の1行目から順に計算していきます．

X = np.zeros((2**N, 26), int)
r = 1
for i in range(N):
    X[r:r*2] = X[:r] + S[i]
    r *= 2

これで，上図でいう表（ここでは行列）になっているので，
最後にK個となる数を数えて最大値をとります．

ans = (X == K).sum(axis=1).max()
print(ans)

以上，ABC249Cの解説でした．

＜参考解説リンク＞

2022-04-17

ABC248C

AtCoder Python

ABC248C の解説

atcoder.jp

解答コード

https://atcoder.jp/contests/abc248/submissions/31033351

問題

数列の和がK以下となるような数列が何通りあるか解く問題です．

解説

C問題にしては，かなーーり難しめな問題2次元DPの問題です．
概要図です．

f:id:nodashin_jpn:20220417082212j:plain — ABC248C_overview

標準入力とMODを定義します．

MOD = 998244353
N, M, K = map(int, input().split())

dpの配列を定義します．
最初の項で取れる場合の数を定義します．

dp = [[0]*(K+1) for _ in range(N)]
dp[0][1: M+1] = [1]*M

順番に項を処理していきます．
第i項をjとしたときの合計が(j+k)となる場合の数を足していきます．
もし(j+k)が与えられたKを超えるようであれば，考える必要はないのでcontinueです．
MODの計算も忘れずにしましょう．

for i in range(1, N):
    for j in range(1, M+1):
        for k in range(1, K+1):
            if k + j > K:
                continue
            dp[i][k+j] += dp[i-1][k]
            dp[i][k+j] %= MOD

最後にK以下の場合をすべて合計して出力します．

ans = sum(dp[-1]) % MOD
print(ans)

以上，ABC248Cの解説でした．

＜参考解説リンク＞

2022-03-13

ABC243D

ABC243D の解説

atcoder.jp

解答コード

https://atcoder.jp/contests/abc243/submissions/30097544

問題

バイナリーツリーの内部処理に必要なグラフの移動を説く教育的な問題です．

解説

移動を示す文字列が，
Uなら「割る2」
Lなら「かける2」
Rなら「かける2 + 1」
と一見シンプルな問題ですが，xの値が非常に大きいことがネックになります．

直近のABCではD問題でも灰-茶であることは珍しくないので，
「これ楽勝じゃね？」
と思ってそのままシミュレートしてTLEくらいました．

時間内に処理を終わらせるには，不要な計算を削減するために無駄な移動を減らす必要があります．

概要図です．

f:id:nodashin_jpn:20220313142647j:plain — ABC243D_overview

では，コードです．まず標準入力を読み取ります．

N, x = map(int, input().split())
S = list(input())

これから，不要な上下運動を削除します．
書き方はいろいろあると思います．例を３つ書きます．
その1）ひとまず s を append して，一番後ろが'U'で，その前が'U'以外なら２個削除する方法．

# その1
T = []
for s in S:
    T.append(s)
    if len(T) > 1:
        if T[-1] == 'U' and T[-2] != 'U':
            T.pop()
            T.pop()

その2）continue を使う．

# その2
T = []
for s in S:
    if s=='U':
        if T:
            if T[-1] != 'U':
                T.pop()
                continue
    T.append(s)

その3）公式解説のように条件文をandでつないで，elseで分岐するなど...

# その3
T = []
for s in S:
    if s=='U' and T and T[-1] != 'U':
        T.pop()
    else:
        T.append(s)

こうして，無駄な動きを排除したTに対して，シミュレートすればOKです．

for t in T:
    if t == 'U':
        x //= 2
    elif t == 'L':
        x = x*2
    elif t == 'R':
        x = x*2 + 1
print(x)

以上，ABC243Dの解法でした．
公式にあるbitの解き方は，まったく発想としてありませんでした．
あの解放が浮かんでくるように精進します...

余談ですが，個人的にこういったアルゴリズムの中身を部分的に抽出した教育的な問題は好きです．
すなわち，コードのコピペで解けるような問題ではなく，きちんとアルゴリズムの中身を理解してるか聞いてくる感じの問題です．そして，解くと勉強になります．
例えば，ワーシャルフロイド法の内容を聞いてくる D - Shortest Path Queries 2 が志向が近い問題です（ただし，ワーシャルフロイド法を使うABC243Eが解けたとは限らないww）．

＜参考解説リンク＞
・公式

問題

解説

具体例 で考える

STEP 1

STEP 2

STEP 3

コード

関連リンク

ABC256F の解説

解答コード

問題

解説

ABC116D の解説

解答コード

問題

解説

ABC250C の解説

解答コード

問題

解説

ABC249C の解説

解答コード

問題

解説

ABC248C の解説

解答コード

問題

解説

ABC243D の解説

解答コード

問題

解説

具体例 $(N, M, K) = (2, 2, 3)$ で考える