🥔

[ABC258E] Packing Potatoes を Pythonで解く

2022/07/03に公開

Python

競技プログラミング

AtCoder

tech

考え方

「尺取り法」と「周期性の処理」の合わせ技です。アルゴリズムは決して高度ではありませんが、実装に苦戦してしたので、備忘録として残しておきます。
　例として、 $N=10$ , $X=10$ のときを考えます。ジャガイモの重さの和が目標 $X=10$ を超えると次の箱になりますが、箱の入れ始めがジャガイモ重量周期のどこから始まっているかで、終わりの位置が決まります。また同時に、箱に入れるジャガイモの個数も決まります。この部分は「尺取り法」を用いて、周期の始まりごとに、終わりの場所・ジャガイモの個数を計算しておきます。後の実装部分では、終わりの場所は Arrow,　ジャガイモの個数は PotatoCntという名前のリストで実装しています。

このジャガイモ重量周期の始まりと終わりを辺で結んだ、グラフに置き換えて考えるのがポイントです。それぞれ出次数が $1$ のグラフになりますが、ジャガイモの箱詰めは、周期 $0$ のところから始まるので、グラフも $0$ から始まります。途中から同じところをまわるので、グラフ頂点の周期を考えれば、計算量を減らすことができます。

実装メモ

実装で苦戦したため、変数名はあえて長めにして意味を取りやすくしています。まず、標準入力の処理、今回はクエリはいったん記録します。

N,Q,X = map(int,input().split())
W = list(map(int,input().split()))

query = []
for _ in range(Q):
    query.append(int(input()))

ジャガイモの一周期分の重量総和を求めます。目標重量 $X$ が大きい場合には、1つの箱に何周期分もジャガイモを入れる必要があるかもしれません。前処理として目標重量 $X$ から周期分を減算しておき、その分最後の解答のときに、一周期あたり $N$ を加えれば良いです。
空箱にあらかじめ $N$ 個、重量SumWgt のジャガイモのセットをいくつか入れて準備しておくとイメージすれば分かりやすいと思います。

SumWgt = sum(W)
Ncnt = 0
while X > SumWgt:
    Ncnt += N * (X//SumWgt)
    X %= SumWgt

尺取り法の部分は、 $l$ , $r$ を決めて、 $l$ をずらしていきます。 $l$ がジャガイモ周期の始まりの場所に相当します。 $l$ をずらしたら目標重量を超えるまで $r$ を右にずらします。 $r$ が終わりの位置で、箱に入るジャガイモの個数は $r-l$ で求められます。

Arrow = [-1] * N
PotatoCnt = [0] * N
s,r = 0,0
for l in range(N):
    if l > 0:
        s -= W[l-1]
    
    while s < X:
        s += W[r%N]
        r += 1

    Arrow[l] = r%N
    PotatoCnt[l] = r-l

グラフ頂点の周期性の処理を行います。周期の大きさ：Cycle、周期性に入る前の頂点数：PreCyeを計算します。これらを計算しておけば、 $k$ 番目の値については、indexを下記に変えたものと同じ値になります。

(k - PreCyc)\ mod\ Cycle + PreCyc

既視のデータは、Seenにindexを入れておきます。見たことのあるデータを見つけた場合には、保存されているindexが、PreCycに相当します。通ってきたグラフの頂点の順序をRecに記録していきます。Rec[0]に記録されているのは、 $1$ 個目の箱のジャガイモ周期の始まりの位置なので、Recの添字は $-1$ しないといけないことに注意です。私はなかなか気付けずデバッグに時間がかかりました。

x = 0
idx = 1
Seen = [-1] * N
Rec = []

while idx :
    if Seen[x] < 0:
        Seen[x] = idx
        Rec.append(x)
        idx += 1
    else:
        PreCyc = Seen[x]
        Cycle = idx - PreCyc
        break

    x = Arrow[x]

クエリで与えられた $k$ がPreCycより後ろであれば、頂点グラフの周期性の処理を行います。前処理でジャガイモを周期ずつまとめて減算していたものを加算します。

ans = []
for k in query:
    if k - PreCyc > 0:
        k = (k - PreCyc) % Cycle + PreCyc
    
    a = PotatoCnt[Rec[k-1]] + Ncnt
    ans.append(a)

print(*ans, sep="\n")

結構大変でした。周期性の処理は苦手です。

考え方

実装メモ

Discussion