Pythonで厳密に増加するカラフルなキャンドル列の数を求めるプログラム

左から右へ一列に並んだ n 本のキャンドルがあるとします。左から i 番目のキャンドルは高さ h[i] と色 c[i] を持っています。さらに整数 k が与えられ、1 から k までの範囲に色が存在することを表します。ここで、「厳密に増加するカラフルなキャンドルのシーケンス」が何通り存在するかを求めるのが問題です。

増加しているかどうかは高さに基づいて判定され、1 から k までのすべての色について、それぞれ少なくとも 1 本のキャンドルが含まれている場合に、そのシーケンスは「カラフル（colorful）」であるとみなされます。答えが非常に大きくなる可能性があるため、結果は 10^9 + 7 で割った余りを返します。

たとえば、入力が K = 3、h = [1,3,2,4]、c = [1,2,2,3] の場合、出力は 2 になります。これは [1,2,4] と [1,3,4] の 2 つのシーケンスが条件を満たすためです。

解法のアプローチ

このアルゴリズムでは、BIT（Binary Indexed Tree／フェニック木）による区間和の高速な取得・更新と、色の部分集合に対する包除原理を組み合わせます。read() と update() はそれぞれBITの累積値参照と更新を行う関数です。具体的には、以下の手順に従います。

• 関数 read() を定義します。引数として T と i を受け取ります。
• s := 0 と初期化します。
• i > 0 の間、次の処理を繰り返します。
  • s := s + T[i]
  • s := s mod 10^9+7
  • i := i − (i AND −i)
• s を返します。
• 関数 update() を定義します。引数として T、i、v を受け取ります。
• i ≤ 50010 の間、次の処理を繰り返します。
  • T[i] := T[i] + v
  • T[i] := T[i] mod 10^9+7
  • i := i + (i AND −i)
• v を返します。
• メイン処理では、以下を実行します。
• L := 2^k、R := 0、N := h のサイズ とします。
• i を 0 から L−1 まで繰り返します。
  • T := サイズ 50009 の配列を作成し、すべて 0 で埋めます。
  • t := 0 とします。
  • j を 0 から N−1 まで繰り返します。
    • (i を (c[j] − 1) 回右シフトした値) が奇数である場合、
      • t := t + update(T, h[j], read(T, h[j] − 1) + 1)
      • t := t mod 10^9+7
  • (i のセットビット数) mod 2 が k mod 2 と等しい場合、
    • R := R + t
    • R := R mod 10^9+7
  • それ以外の場合、
    • R := (R + 10^9+7) − t
    • R := R mod 10^9+7
• R を返します。

実装例

理解を深めるために、以下の実装を見てみましょう。

def solve(k, h, c):
    def read(T, i):
        s = 0
        while i > 0:
            s += T[i]
            s %= 1000000007
            i -= (i & -i)
        return s

    def update(T, i, v):
        while i <= 50010:
            T[i] += v
            T[i] %= 1000000007
            i += (i & -i)
        return v

    def number_of_bits(b):
        c = 0
        while b:
            b &= b - 1
            c += 1
        return c

    L = 2 ** k
    R = 0
    N = len(h)

    for i in range(L):
        T = [0 for _ in range(50010)]
        t = 0

        for j in range(N):
            if (i >> (c[j] - 1)) & 1:
                t += update(T, h[j], read(T, h[j] - 1) + 1)
                t %= 1000000007

        if number_of_bits(i) % 2 == k % 2:
            R += t
            R %= 1000000007
        else:
            R += 1000000007 - t
            R %= 1000000007
    return R

k = 3
h = [1,3,2,4]
c = [1,2,2,3]

print(solve(k, h, c))

入力

3, [1,3,2,4], [1,2,2,3]

出力

2