条件を満たす(i, j) の数を数える問題

どういうことか

問題
整数列a = { $a _ 0, a _ 1, \dots , a _ {n-1}$ } が与えられる。a[i] = a[j] ( i < j ) となるような(i, j) の組の数を求めてください。 $( n \leq 10 ^ 5 )$

(i, j) の組を全探索すると当然求められるが、 $O(n ^ 2 )$ でTLEしてしまう。

そこで、片方を固定することを考えてみる。j を固定して、i < j を満たすi の中でa[j] = a[i] を満たすものがいくつあるかを数えていくことにする。これは連想配列を用いると簡単である。C++ならstd::map。j を昇順に見ていき、今までに出てきたa[i] は連想配列に記録しておく。そんな感じで $O( N log (N) )$ で解ける。

以下実装。組の数は最大n*(n+1)/2 なので32 bit 整数型に収まらない可能性があり注意。

#include <iostream>
#include <algorithm>
#include <vector>
#include <map>
#define rep(i,n) for(int i=0;i<n;++i)
#define rep1(i,n) for(int i=1;i<=n;++i)
using namespace std;
template<class T>bool chmax(T &a, const T &b) { if(a < b){ a = b; return 1; } return 0; }
template<class T>bool chmin(T &a, const T &b) { if(a > b){ a = b; return 1; } return 0; }
typedef long long ll;
int main()
{
  int n; cin >> n;
  vector<int> a(n);
  rep(i,n) cin >> a[i];

  map<int,ll> mp;
  ll res = 0;
  rep(i,n) {
    res += mp[a[i]];
    mp[a[i]]++;
  }
  cout << res << "\n";
  
  return 0;
}

例題

a[i] = a[j] を満たす(i, j) の数を求める。という問題に帰着する問題はよくある。

条件はA(i) = A(j) でなく、A(i) = B(j) でも良い

・E - This Message Will Self-Destruct in 5s

問題の条件を式にすると、 $j - i = a(i) + a(j) \; ( j > i )$ となる。
これは簡単な式変形により、 $a(i) + i = j - a(j) \; ( j > i )$ となり、A(i) = a(i) + i , B(i) = i - a(i) とすれば、 $A(i) = B(j) \; ( j > i )$ を満たす(i, j) を求めれば良い問題になる。

累積和を使う

・A - Zero-Sum Ranges

悪名高き最難200点問題。連続する部分列の総和を累積和 $S(x) := \sum _ {i=0} ^ {x-1} a(i)$ を用いて表すと、 $a(l) + a(l+1) + \dots + a(r-1) = S(r) - S(l)$ 。これを用いると問題は $S(r) - S(l) = 0 \; ( r > l )$ を満たす(l, r) の組を求める問題に言い換えられる。これは最初の問題と同じ問題である。

連想配列の更新を少し工夫する

・E - Rem of Sum is Num

問題の条件は、累積和を用いて $( S(r) - S(l) ) \% K = r - l \; ( r > l )$ 。これは言い換えると、 $S(r) - r = S(l) - l \; ( mod K ) ( K > r - l )$ となる。
K > r - l という条件があるので、連想配列に値を記録するだけでなく削除していけば良い。

ちょっと変わった累積和

・D - Multiple of 2019

桁DPしたくなる、けど間に合わない（実は間に合う）。
文字列の長さをn とする。x 文字目からy 文字目を10進数としてみた値をS(x, y) とすると、 $S(i, j) = \frac{S(i, n) - S(j+1, n)}{10 ^ {n - j} }$ となる。ここで2019 は2 でも5 でも割り切れないので、 $S(i, n) - S(j+1, n) = S(i, j) * 10 ^ {n - j} = S(i, j) \; (mod 2019 )$ となる。ので、結局Zero-Sum Ranges と同じで、 $S(i, n) = S(j+1, n)$ を満たす(i, j) を数えれば良い。