目次
はじめに
この記事では、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です。
例
ランレングス圧縮とその復元のコードは手元に持っておくとよいでしょう。
コード
from itertools import groupby
def runLengthEncode(S):
res = []
for k, v in groupby(S):
count = sum(1 for _ in v)
res.append([k, count])
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問題の解説でした！
では、またね。
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”