Pythonでランレングス符号化ベクトルの内積を求めるプログラム

問題概要

ランレングス符号化(Run-Length Encoding / RLE)は、連続して現れる同じ値を「個数」と「値」のペアで表現するコンパクトなデータ形式です。例えば、ベクトル [1, 1, 1, 2, 2, 2, 2] は「1が3個続き、その後に2が4個続く」という意味で [3, 1, 4, 2] と表されます。

本記事では、この形式で与えられた2つのベクトル nums1 と nums2 の内積(ドット積)を求めるPythonプログラムを紹介します。内積とは、2つのベクトルの対応する要素同士を掛け合わせ、その総和を取った値のことです。

具体例で確認

入力が nums1 = [2, 7, 5, 3]、nums2 = [3, 5, 4, 2] の場合、これらはそれぞれ次のようなベクトルを表しています。

• nums1 → [7, 7, 3, 3, 3, 3, 3](7が2個、3が5個)
• nums2 → [5, 5, 5, 2, 2, 2, 2](5が3個、2が4個)

したがって、内積は次のように計算されます。

[7, 7, 3, 3, 3, 3, 3] ・ [5, 5, 5, 2, 2, 2, 2]
= 7×5 + 7×5 + 3×5 + 3×2 + 3×2 + 3×2 + 3×2
= 35 + 35 + 15 + 6 + 6 + 6 + 6
= 109

アルゴリズムの考え方

ベクトルをいったん展開してから内積を計算しても正しい結果は得られますが、要素数が多い場合はメモリと計算時間を大きく浪費してしまいます。そこで、RLE形式のまま各ブロック(個数と値のペア)を末尾から順に取り出し、両ベクトルで重なり合う区間だけを掛けて加算していくのが効率的です。片方のブロックに余りが生じた場合は、その余りを元のリストに戻し、次のループで再利用します。

1. 結果を格納する変数 ans を 0 で初期化します。
2. nums1 と nums2 がどちらも空でない間、以下の処理を繰り返します。
   • nums1 の末尾から値 val1 と個数 count1 を取り出します(pop)
   • nums2 の末尾から値 val2 と個数 count2 を取り出します(pop)
   • ans に (val1 × val2) × min(count1, count2) を加算します
   • count2 > count1 の場合:余りの個数 |count2 − count1| と値 val2 を nums2 の末尾に戻します
   • count1 > count2 の場合:余りの個数 |count2 − count1| と値 val1 を nums1 の末尾に戻します
3. すべてのブロックを処理し終えたら、ans を返します。

Pythonでの実装例

以下のコードを実行すると、期待どおりの結果が得られることを確認できます。

def solve(nums1, nums2):
    ans = 0

    while nums1 and nums2:
        val1 = nums1.pop()
        count1 = nums1.pop()
        val2 = nums2.pop()
        count2 = nums2.pop()

        ans += (val1 * val2) * min(count2, count1)

        if count2 > count1:
            nums2.append(abs(count2 - count1))
            nums2.append(val2)
        elif count1 > count2:
            nums1.append(abs(count2 - count1))
            nums1.append(val1)

    return ans

nums1 = [2, 7, 5, 3]
nums2 = [3, 5, 4, 2]
print(solve(nums1, nums2))

入力

[2, 7, 5, 3], [3, 5, 4, 2]

出力

109

計算量の目安

このアルゴリズムの時間計算量は O(m + n) です(m・n はそれぞれのベクトルのブロック数)。入力リストを直接操作するため、追加のメモリもほとんど必要ありません。ベクトルを展開してから計算する方法(全要素数に比例した時間とメモリが必要)と比べて効率が良く、特に長いランを含むデータで威力を発揮します。