Pythonでリストのすべての部分列の総和Sに対する2^Sの合計を効率的に求めるプログラム

リスト A が与えられたとします。ここで、A のすべての空でない部分列(サブリスト)を考えます。n 個の要素を持つリストには (2ⁿ − 1) 個の空でない部分列が存在することが知られています。それぞれの部分列について要素の総和(sublist_sum)を計算し、それらを S₁, S₂, S₃, …, S_(2N−1) と表します。

そして、次のような特別な総和 P を定義します。

P = 2^S1 + 2^S2 + 2^S3 + … + 2^S(2N−1)

この P の値を求めるのが目的です。ただし、P は非常に大きな値になる可能性があるため、P mod (10⁹ + 7) を返します。

入力例と出力例

たとえば、入力が A = [2, 2, 3] の場合、出力は 224 になります。各部分列と対応する値は次のとおりです。

• {2} → 2² = 4
• {2} → 2² = 4
• {3} → 2³ = 8
• {2, 2} → 2⁴ = 16
• {2, 3} → 2⁵ = 32
• {2, 3} → 2⁵ = 32
• {2, 2, 3} → 2⁷ = 128

これらを合計すると、4 + 4 + 8 + 16 + 32 + 32 + 128 = 224 となります。

解法のアプローチ

すべての部分列を実際に生成していては、部分列の数が要素数に対して指数的に増えるため非効率です。ここで重要なのは、各要素について「その要素を部分列に含めるか、含めないか」という2択しかないという点です。積 ∏(1 + 2^el) を展開すると、各項がちょうど1つの部分列に対応するため、次の関係が成り立ちます。

空でない部分列に対する P = ∏(1 + 2^el) − 1

「−1」は、空集合(総和 0、寄与 2⁰ = 1)の分を差し引くためです。したがって、以下の手順で O(N) 時間で答えを求められます。

• ans := 1、m := 10⁹ + 7 と初期化する
• A の各要素 el について:
  • ans := ans × (1 + (2^el mod m))
  • ans := ans mod m
• (m + ans − 1) mod m を返す

Pythonでの実装例

以下の実装を見ると、仕組みの理解が深まるでしょう。

def solve(A):
   ans=1
   m=10**9+7

   for el in A:
      ans *= (1+pow(2,el,m))
      ans %= m
   return (m+ans-1) % m


A = [2,2,3]
print(solve(A))

入力

[2,2,3]

出力

224