Pythonで文字列を書き込む行数と最終行の幅を求めるアルゴリズム

問題の概要

文字列 S が与えられたとき、その文字を左から右へ順番に行へ書き込んでいくことを考えます。各行の最大幅は 100 単位で、ある文字を書き込んだ結果、行の幅が 100 単位を超えてしまう場合は、その文字は次の行に書き込まれます。さらに、配列 widths も与えられており、widths[0] は「a」の幅、widths[1] は「b」の幅、といった具合に各アルファベットの幅が格納されています。

この問題では、次の 2 つの問いに答える必要があります。

• 文字列 S の少なくとも 1 文字が書き込まれている行は何行あるか？
• 最後の行で使用されている幅は何単位か？

答えは、長さ 2 の整数リストとして返します。

具体例で確認

たとえば、次の入力を考えてみましょう。

• widths = [4,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10]
• S = "bbbcccdddaaa"

この場合の出力は [2, 4] になります。「a」以外のすべての文字は幅が 10 で共通しており、"bbbcccdddaa" までで 9 × 10 + 2 × 4 = 98 単位を使用します。最後の「a」については、1 行目には残り 2 単位しかないため、2 行目に書き込まれます。したがって、答えは「全体で 2 行」、そして「2 行目の幅は 4 単位」となります。

解き方のアプローチ

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

1. line を 1、count を 0 で初期化します。
2. 文字列 S の各文字 i について、次の処理を繰り返します。
   • count に widths[i の ASCII コード − 97] を加算します。
   • count が 100 を超えた場合は、line を 1 増やし、count を widths[i の ASCII コード − 97] にリセットします。
3. 最後に [line, count] を返します。

実装例

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

class Solution:
    def numberOfLines(self, widths, s):
        line = 1
        count = 0
        for c in s:
            count += widths[ord(c) - 97]
            if count > 100:
                line += 1
                count = widths[ord(c) - 97]
        return [line, count]

ob = Solution()
print(ob.numberOfLines([4,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10], "bbbcccdddaaa"))

入力

[4,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10],"bbbcccdddaaa"

出力

[2, 4]

コードのポイント

このアルゴリズムの時間計算量は O(n)（n は文字列の長さ）、追加のメモリ使用量は O(1) と非常に効率的です。ord(c) - 97 によって、各文字を配列 widths のインデックス（a = 0、b = 1、…）に変換できる点がポイントです。また、行の折り返し判定は「現在の count に新しい文字の幅を加算して 100 を超えるかどうか」だけで行えるため、文字列を一度走査するだけで答えを求められます。