問題

$n$ 個の整数からなる数列 $a$ が与えられる。

この数列の要素 ${a}_{i}$ に対し、 ${b}_{i} = {2}^{{a}_{i}}$ と定めるとき、 ${{b}_{i}}^{{b}_{j}} = {{b}_{j}}^{{b}_{i}}$ が成立するような整数 $(i, j)$ の組（ただし、 $1 \leq i \lt j \leq n$ ）の個数を求めよ。

入力

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

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

最初の1行目には、整数 $n$ が与えられる。
次の1行には、 $n$ 個の整数 ${a}_{i}$ が順に与えられる。

条件

$1 \leq t \leq {10}^{4}$
- すべてのテストケースにおける $n$ の値の合計は $2 \cdot {10}^{5}$ を超えない。
$1 \leq {a}_{i} \leq {10}^{9}$

出力

$t$ 行にわたり、各テストケースでの答えを1行ごとに出力すること。

解法

方程式 ${{b}_{i}}^{{b}_{j}} = {{b}_{j}}^{{b}_{i}}$ が、どのようなときに成立するのかについて考えます。

まず、これは ${b}_{i} = {b}_{j}$ のときには必ず成立します。すなわち、 ${a}_{i} = {a}_{j}$ であるような $(i, j)$ の組はすべて条件を満たすと言えます。

では、 ${a}_{i} \ne {a}_{j}$ であるような $(i, j)$ に対して、元の方程式が成立するような例はあるのかということを考えます。そのために与えられた方程式から、 \begin{align} {({2}^{{a}_{i}})}^{{b}_{j}} & = {({2}^{{a}_{j}})}^{{b}_{i}} \\ {2}^{{a}_{i} \cdot {b}_{j}} & = {2}^{{a}_{j} \cdot {b}_{i}} \\ {a}_{i} \cdot {b}_{j} & = {a}_{j} \cdot {b}_{i} \\ {a}_{i} \cdot {2}^{{a}_{j}} & = {a}_{j} \cdot {2}^{{a}_{i}} \\ \frac{{2}^{{a}_{i}}}{{a}_{i}} & = \frac{{2}^{{a}_{j}}}{{a}_{j}} \end{align} と整理することができます。

この整理から、 $x \gt 0$ に対して関数 $f(x)$ を $f(x) = \displaystyle \frac{{2}^{x}}{x}$ と定義します。このとき、 \begin{align} f'(x) & = \left( - \frac{1}{{x}^{2}} \right) \cdot {2}^{x} + \frac{1}{x} \cdot \left( {2}^{x} \log 2 \right) \\ & = \frac{{2}^{x}}{{x}^{2}} \left( (\log 2) x - 1 \right) \end{align} と整理できます。

ここで、 $f'(x) = 0$ となる $x$ を $\alpha$ と定義します。すなわち、 $(\log 2) \alpha - 1 = 0$ とします。

$(\log 2) x - 1$ の値が単調増加であることから、 $0 \lt x \leq \alpha$ の範囲で $f'(x) \leq 0$ であり、 $\alpha \leq x$ の範囲で $f'(x) \geq 0$ が言えます。

従って、 $f(x)$ は $0 \lt x \leq \alpha$ で単調減少であり、 $\alpha \leq x$ で単調増加であると言えます。

ここで、 $2 \lt e \lt {2}^{2}$ より、 $\log 2 \lt 1 \lt 2 \log 2$ が成立することから、 $1 \lt \alpha \lt 2$ です。

元の方程式より、実際に考える必要があるのは $x$ が $1$ 以上の整数のときについてなので、 $f(x)$ は $x \geq 2$ のときに単調増加であると言えます。

よって ${a}_{i}, {a}_{j} \geq 2$ のときについては、 ${a}_{i} \ne {a}_{j}$ であれば、元の方程式が成立するような $(i, j)$ の組み合わせは存在しないと言えます。

次に、 $f(x) = f(1)$ となるような整数 $x$ （ただし、 $x \ne 1$ ）が存在するかについて考えます。ここで、 \begin{align} f(1) = \frac{{2}^{1}}{1} = 2 \end{align} です。

$x \geq 2$ については $f(x)$ は狭義単調増加であることから、もし $x$ を小さい順に探索していき、 $f(x) = 2$ であるような整数 $x$ が存在すれば、その $x$ が唯一の $x \geq 2$ での解と言えます。また、同様の探索方法で $f(x) \gt 2$ であるような $x$ が存在したら、そのような整数 $x$ は存在しないと言えます。

実際には、 \begin{align} f(2) = \frac{{2}^{2}}{2} = 2 \end{align} であるので、 $x = 2$ が条件を満たす整数 $x$ であると言えます。

以上の考察から、元の方程式を満たすような $({a}_{i}, {a}_{j})$ の値の条件は

${a}_{i} = {a}_{j}$ であること
$({a}_{i}, {a}_{j}) = (1, 2)$ もしくは $(2, 1)$ であること

のいずれかとなります。

ここで、問題文の条件の通りに $1 \leq i \lt j \leq n$ で考えられる $(i, j)$ に対して ${a}_{i}, {a}_{j}$ が条件を満たすかを判定すると、 $O( {n}^{2})$ の計算量がかかってしまい、実行時間制限をオーバーしてしまいます。

そこで、 ${a}_{i} = p$ である $i$ の個数を $\mathrm{cnt}_{p}$ として定義します。このとき、 ${a}_{i} = {a}_{j} = p$ である $(i, j)$ （ただし $i \lt j$ ）の個数は \begin{align} \frac{1}{2} \mathrm{cnt}_{p} (\mathrm{cnt}_{p} - 1) \end{align} と計算することができます。

また、 $({a}_{i}, {a}_{j}) = (1, 2)$ もしくは $(2, 1)$ である $(i, j)$ の個数は \begin{align} \mathrm{cnt}_{1} \cdot \mathrm{cnt}_{2} \end{align} によって計算することができます。

従って、数列 $a$ のすべての要素を見て $\mathrm{cnt}$ を用意しておき、 \begin{align} \sum_{p \in a} \frac{1}{2} \mathrm{cnt}_{p} (\mathrm{cnt}_{p} - 1) + \mathrm{cnt}_{1} \cdot \mathrm{cnt}_{2} \end{align} の値が求める答えとなります。

これは、 $O(n)$ の計算量で求めることができるので、実行時間制限にも十分間に合います。

ソースコード

答えとなる組み合わせの個数を計算する solve() を以下のように実装しました。

int n, a[int(2e5 + 5)];
map<int, ll> cnt; // a[i] は最大で 10^9 になりえるので、map を用いた

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

  ll ans = cnt[1] * cnt[2]; // ans は答えとなる値であり、予め (1, 2) の組み合わせの総数を計算しておく
  for (int i = 1; i <= n; i++) {
    ans += cnt[a[i]] * (cnt[a[i]] - 1) / 2; // cnt[a[i]] = 0, 1 の場合はこの値は 0 になるので特に問題ない
    cnt[a[i]] = 0; // 一旦計算した値は 0 にしておくことで、重複しての計算を防げる
  }

  return ans;
}

提出したコード: https://codeforces.com/contest/1899/submission/237871519
GitHub

感想

問題を解いているときは、問題中の入力 / 出力例でなんとなく $(1, 2)$ の組み合わせ以外は存在しなさそうだなと思いました。実際に記事として書くために証明するとなると結構時間がかかりました。

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

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

D - Yarik and Musical Notes / Codeforces Round 909 (Div. 3)

問題

入力

条件

出力

解法

ソースコード

感想