Pythonでk個の連続する重複文字を削除した後の文字列を求めるプログラム

文字列 s と整数 k が与えられたとき、「同じ文字が k 個連続している部分」を繰り返し削除していき、最終的に残る文字列を求める問題を考えます。

例えば、入力が s = "paaappmmmma"、k = 3 の場合、出力は "ma" になります。処理の流れは以下のとおりです。

• まず連続する3つの "a" を削除 → "pppmmmma"
• 次に連続する3つの "p" を削除 → "mmmma"
• 最後に4つある "m" のうち連続する3つを削除 → "ma"

解き方のアプローチ

この問題は、次の手順で解くことができます。

• 以下の処理を、変更がなくなるまで繰り返します。
  • count を 0 で初期化する
  • s に含まれる一意な文字の集合 chars を取得する
  • chars の各文字 c について、次を行う
    • c が k 個連続した部分文字列が s に存在する場合は、それを削除し、count を 1 増やす
  • count が 0 のままなら、これ以上削除できる箇所がないためループを抜ける
• 最終的な s を返す

実装例

理解を深めるために、以下のPythonコードを見てみましょう。

class Solution:
    def solve(self, s, k):
        while True:
            count = 0
            chars = set(s)
            for c in chars:
                if c * k in s:
                    s = s.replace(c * k, "")
                    count += 1
            if count == 0:
                break
        return s

ob = Solution()
s = "paaappmmmma"
k = 3
print(ob.solve(s, k))

入力

"paaappmmmma", 3

出力

ma

計算量の改善：スタックを使った O(n) の解法

上記の素朴な実装では、削除がなくなるまで文字列全体を何度も走査するため、文字列が長い場合には非効率になる可能性があります。スタックを使えば、文字列を一度だけ走査する O(n) の計算量で同じ結果を得られます。

class Solution:
    def solve(self, s, k):
        stack = []  # [文字, 連続数] のペアを管理
        for ch in s:
            if stack and stack[-1][0] == ch:
                stack[-1][1] += 1
                if stack[-1][1] == k:
                    stack.pop()
            else:
                stack.append([ch, 1])
        return ''.join(ch * cnt for ch, cnt in stack)

ob = Solution()
s = "paaappmmmma"
k = 3
print(ob.solve(s, k))

この方法では、各文字について「直前の文字と同じかどうか」と「連続数が k に達したか」を確認するだけでよいため、大きな入力に対しても高速に動作します。連続数が k に達した時点でそのグループをスタックから取り除くことで、削除をその場で完了できるのがポイントです。