ABC380 C問題の解説(Python)

はじめに

この記事では、ABC380のC問題を解説していきます。

ABC(AtCoder Beginner Contest)とは、AtCoderが開催している、競技プログラミングコンテストです。

ABC380 C – Move Segment

問題

問題

0, 1 のみからなる長さ N の文字列 S が与えられます。
S の中で先頭から K 番目の 1 の塊を K-1 番目の 1 の塊の直後まで移動した文字列を出力してください。

なお、S には 1 の塊が K 個以上含まれることが保証されます。

制約

  • \(1 \leqq N \leqq 5×10^5\)
  • \(2 \leqq K\)

思考の筋道

同じ文字が並んでいる箇所を塊として考えたいので、ランレングス圧縮を使いましょう。

ランレングス圧縮とは、”001110″ のような文字列を [(“0”, 2), (“1”, 3), (“0”, 1)] と変換することを言います。(0が2個並んで、1が3個並んで、0が1個並ぶ)

今回の問題では、与えられた文字列をランレングス圧縮し、K番目の1の塊を直前の塊と入れ替え、その後に文字列に復元すればOKです。

S = “010011100011001”, K = 3 とします。

ランレングス圧縮すると次のようになります。
[(“0”, 1), (“1”, 1), (“0”, 2), (“1”, 3), (“0”, 3), (“1”, 2), (“0”, 2), (“1”, 1)]

3度目の1の塊を直前の塊と入れ替えます。
[(“0”, 1), (“1”, 1), (“0”, 2), (“1”, 3), (“1”, 2), (“0”, 3), (“0”, 2), (“1”, 1)]

文字列に復元します。
“010011111000001”

ランレングス圧縮とその復元のコードは手元に持っておくとよいでしょう。

コード

from itertools import groupby


def runLengthEncode(S):
    res = []
    for k, v in groupby(S):
        res.append((k, len(list(v))))
    return res


def runLengthDecode(L):
    res_list = []
    for c, n in L:
        res_list.append(c * int(n))
    return "".join(res_list)


N, K = map(int, input().split())
S = input()

l = runLengthEncode(S)
count = 0
for i in range(len(l)):
    if l[i][0] == "1":
        count += 1
    if count == K:
        l[i], l[i - 1] = l[i - 1], l[i]

print(runLengthDecode(l))

ちなみに、私がコンテスト中に問題を解いた時は、ランレングス圧縮を使わずにもっと原始的に解いていました。供養として載せておきます。

N, K = map(int, input().split())
S = input()

count = 0
end_K_1 = 0
start_K = 0
end_K = 0

for i in range(N):
    if S[i] == "1":
        if i == 0 or S[i - 1] == "0":
            count += 1
            if count == K:
                start_K = i
        if i == N - 1 or S[i + 1] == "0":
            if count == K - 1:
                end_K_1 = i
            elif count == K:
                end_K = i

print(S[0 : end_K_1 + 1] + S[start_K : end_K + 1] + S[end_K_1 + 1 : start_K] + S[end_K + 1 : N])

以上、ABC380のC問題の解説でした!

では、またね。

リンク