D - Distinct Trio Editorial by Nyaan


ユーザー解説です

この問題は DP (動的計画法) で解くこともできます。

この問題は \(A\) から異なる値を \(3\) 個選んで集合を作る方法の通り数を数える問題です。そこで、\(1,2,\dots\) の順に 「\(n\) を集合に入れるか?」を考えていきましょう。
DP テーブルを次のように定義します。

  • \(dp_{i,j} :=\) \(i\) まで確定していて、集合のサイズが \(j\) である通り数。( \(i,j\)\(0 \leq i \leq \max(A), 0 \leq j \leq 3\) の範囲)

DP の遷移を考えてみましょう。集合に \(i\) を入れる場合とそうでない場合で通り数は次のように変化します。

  • 集合に \(i\) を追加する場合: \(i\) の選び方は \(A\) に含まれる \(i\) の個数の分だけ存在する。また集合のサイズは \(1\) 増える。
  • 集合に \(i\) を追加しない場合:通り数・集合のサイズはそのままである。

よって、次の式によって DP テーブルを計算していくことができます。

  • \(dp_{i+1,j} = dp_{i,j} + dp_{i,j-1} \times (A に含まれる i の個数)\)

また、求めたい答えは \(dp_{\max(A),3}\) です。よって DP を計算することでこの問題を解くことができます。計算量は \(\mathrm{O}(N + \max(A))\) です。

Python による実装例は次の通りです。簡単のため DP は 1 次元テーブルを使い回す方法で実装しています。

N = int(input())
freq = [0] * (2 * 10**5 + 10)
for a in map(int, input().split()):
  freq[a] += 1
dp = [1, 0, 0, 0, 0]
for f in freq:
  nx = [0] * 5
  for i in range(4):
    nx[i + 1] += dp[i] * f
    nx[i] += dp[i]
  dp = nx
print(dp[3])

posted:
last update: