yutasの競技プログラミング勉強帖

競技プログラミングの問題についての解説記事を主に書いています。

B - Forming Triangles / Educational Codeforces Round 161 (Rated for Div. 2)

CF1922 Codeforces C++ Div.2 Educational Codeforces Round 161 Diff 1200 幾何問題 競技プログラミング

問題

n 本の棒があり、 i 番目の棒は {2}^{{a}_{i}} の長さである。

ここから 3 本の棒を選び三角形を作るとき、 3 本の棒の選び方は何通りあるか求めよ。

入力

まず最初の1行目に、テストケースの個数を表す整数 t が与えられる。

その後、 t 個のテストケースのそれぞれに対して、以下のように入力が与えられる。

  • 最初の1行目には、整数 n が与えられる。
  • 次の1行には、 n 個の整数 {a}_{1}, \, {a}_{2}, \, \cdots , \, {a}_{n} が順に与えられる。

条件

  • 実行時間制限: 2s
  • メモリ制限: 256MB
  • 1 \leq t \leq {10}^{4}
  • 1 \leq n \leq 3 \cdot {10}^{5}
    • すべてのテストケースにおける n の値の合計は 3 \cdot {10}^{5} を超えない。
  • 0 \leq {a}_{i} \leq n

出力

t 行に渡り、各テストケースについて答えを1行ずつ出力すること。

解法

選んだ3本の棒の {a}_{i} の値を x \leq y \leq z とします。 このとき、長さ {2}^{x}, \, {2}^{y}, \, {2}^{z} の棒で三角形を作ることができる条件について考えます。

まず、 2 \cdot {2}^{z - 1} = {2}^{z} であることから、長さ {2}^{z - 1}, \, {2}^{z - 1}, \, {2}^{z} の3本の棒では三角形は作成することができません。 （2本の長さ {2}^{z - 1} の棒をまっすぐ並べることにより、長さ {2}^{z} の棒と同じ長さになるためです。）

従って、 x, \, y, \, z は整数であるので、 y \lt z の場合では三角形を作ることはできません。

以上から、 y = z の場合について考えます。 この場合では、 x \leq y = z であるので、 y = z の二等辺三角形または正三角形を作成することができます。

このことから、 x \leq y = z となるような3本の棒の組み合わせの個数の総数を求めます。

ここで、 {a}_{i} = j となる i の個数を \mathrm{cnt}_{j} とします。 また、 {a}_{i} \lt j となる i の個数を {s}_{j} とします。

このとき、等しい辺の長さが {2}^{j} である二等辺三角形（正三角形ではない）の個数は、 \begin{align} \frac{\mathrm{cnt}_{j} \cdot (\mathrm{cnt}_{j} - 1)}{2!} \cdot {s}_{j} \end{align} で求めることができます。 また、1辺の長さが {2}^{j} である正三角形の個数は、 \begin{align} \frac{\mathrm{cnt}_{j} \cdot (\mathrm{cnt}_{j} - 1) \cdot (\mathrm{cnt}_{j} - 2)}{3!} \end{align} で求めることができます。

従って、最も長い辺の長さが {2}^{j} である三角形の個数は、 \begin{align} \frac{\mathrm{cnt}_{j} (\mathrm{cnt}_{j} - 1)}{2} \cdot {s}_{j} + \frac{\mathrm{cnt}_{j} (\mathrm{cnt}_{j} - 1) (\mathrm{cnt}_{j} - 2)}{6} \end{align} で求めることができます。

よって、 0 \leq {a}_{i} \leq n であることから、三角形の総数は、 \begin{align} \sum_{j = 0}^{n} \left( \frac{\mathrm{cnt}_{j} (\mathrm{cnt}_{j} - 1)}{2} \cdot {s}_{j} + \frac{\mathrm{cnt}_{j} (\mathrm{cnt}_{j} - 1) (\mathrm{cnt}_{j} - 2)}{6} \right) \end{align} で与えられます。

この方法では O(n) の計算量で求めることができるため、これは実行時間制限に十分間に合うと言えます。

ソースコード

各テストケースに対して、作成できる三角形の個数を返す関数 solve() を以下のように実装しました。

int n;
ll a[int(3e5 + 5)], cnt[int(3e5 + 5)];

ll solve() {
  // 値の入力
  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> a[i];
    cnt[a[i]]++; // cnt の値を更新しておく
  }

  ll ans = 0, sum = 0; // ans は最終的に返す値を、 sum は記事中の s を表す

  // i = 0 から n の順に、最長の辺の長さが 2^i である三角形の個数を計算していく
  for (int i = 0; i <= n; i++) {
    ans += cnt[i] * (cnt[i] - 1) * (cnt[i] - 2) / 6; // 正三角形の個数の更新
    ans += cnt[i] * (cnt[i] - 1) * sum / 2; // 二等辺三角形の個数の更新
    sum += cnt[i]; // 次の i のために、 sum の値を更新する
  }

  return ans;
}