Pythonで最初と最後の要素が同じサブリストの個数を求めるプログラム

数値のリスト nums が与えられたとき、最初の要素と最後の要素が一致するサブリスト（部分リスト）の個数を求めることを考えます。

たとえば、入力が nums = [10, 15, 13, 10] の場合、答えは 5 になります。条件を満たすサブリストは次の5つです。

• [10]
• [15]
• [13]
• [10]
• [10, 15, 13, 10]

解法のアプローチ

この問題は、各要素の出現回数を数えて組み合わせの公式を適用することで、O(n) の計算量で効率的に解けます。手順は以下の通りです。

• 単一要素のサブリストは必ず条件を満たすため、初期値として num_sublists := len(nums) を設定します。
• d := 空の辞書（マップ） を用意します。
• nums 内の各要素 n について、d[n] := d[n] + 1 として出現回数をカウントします。
• 辞書 d 内の各値 k とその出現回数 v について、v が 1 でない場合、num_sublists := num_sublists + (v−1) × v ÷ 2 を加算します。
• 最後に num_sublists を返します。

なぜこの式で求まるのか？

同じ値が v 個あるとき、その中から異なる2つの位置を選ぶ方法は、組み合わせの公式 C(v, 2) = v × (v − 1) / 2 通りあります。選んだ2つの位置をサブリストの始点と終点にすれば、そのサブリストの最初と最後の要素は必ず一致します。これに単一要素のサブリスト分 len(nums) を加えることで、答えが得られます。

実装例

from collections import defaultdict
class Solution:
    def solve(self, nums):
        # 単一要素のサブリストは常に条件を満たす
        num_sublists = len(nums)

        # 各要素の出現回数をカウント
        d = defaultdict(int)
        for n in nums:
            d[n] += 1

        # 同じ値のペアの組み合わせ数を加算
        for k, v in d.items():
            if v != 1:
                num_sublists += (v - 1) * v // 2
        return num_sublists

ob = Solution()
nums = [10, 15, 13, 10]
print(ob.solve(nums))

入力

[10, 15, 13, 10]

出力

5

計算量

時間計算量・空間計算量はともに O(n) です。すべてのサブリストを列挙して確認する素朴な O(n²) のアプローチと比べて、大幅に高速に動作します。