ABC206 C 解答

C - Swappable（171）

問題

C - Swappable

問題文

N 個の整数からなる配列 A = (A1, A2, ... , AN)
が与えられるので、次の条件を全て満たす整数組 (i, j)の数を求めてください。
・1 ≤ i < j ≤ N
・Ai≠Aj

制約

入力は全て整数
2 ≤ N ≤ 3×10^5
1 ≤ Ai ≤ 10^9

考察

異なる数字の組み合わせを数え上げます。本問題を特に工夫せずに解こうとすると、計算量はO(N^2)です。2 ≤ N ≤ 3×10^5ですので、ちょっと間に合わないです。

ということで、条件を一つずつ見ていき一歩ずつ高速化していきましょう。その際に次の例を用いて考えていきます。

1 2 3 2 4 5 6 2 5 3

条件は２つありますね。

・1 ≤ i < j ≤ N
・Ai≠Aj

これ、上の条件だけだったら簡単に求められますね。自分より右側に数字が何個あるかを足していけば答えが求まります。

具体的にはこんな感じです。

これならfor文を1回実行すれば答えが求まります。

この考えに下の条件を追加してみます。下の条件では２を考えたときにダメとなる場合が起こりそうです。１の隣の２は、上の条件だけなら8個組み合わせが考えられましたが、それよりも右に2個同じ数字が出てきています。こいつらを除かなければなりません。

8個から2個を除いた値がこの場合の組み合わせ数となります。

逆に言いますと、これを除いてあげるだけで答えが求まります。

ということで、それに必要な値を先に計算してあげましょう。今回は数列Aに各数字が何回出現したかを用いてあげます。これに加えて、先頭から計算していくときにその時点で各数字が何回出現したか（自分を含め左に何回出現したか）を計算しておくことで、あっという間に自分よりも右に同じ数字が何回出てきたかを求まることができます。

図で示すと次のようになります。

先頭の１は重複しないのであまり面白くないですね。素直に、右にある数字9を答えに足しましょう。

ちょっと飛ばして２個目の２です。全体の2の数３からこれまでに出現した２の数２を引くことで、右には２は１個あることがわかりました。この2より右側には数字は６個ありますので、６－１＝５を答えに足してあげましょう。

これを、全部の数字に対して行えば答えが求まります。前計算にてfor文を使っていますが、for文の中でfor文を回しているわけではないので、1重ループのO(N)で答えを求めることができました。

最後にちょっとだけ注意です。制約を見ますと、Aの値はとっても大きいですね。ですので、出現数をカウントするときにAの最大値までの配列を用意してしまうとメモリが足りなくなってしまいます。ですので、今回はmapというデータ構造を用いています。いわゆる辞書型というやつで必要に応じてその値の場所を確保してくれる優れものです。

実装

#include<bits/stdc++.h>
#define rep(i,n) for(int i=0;i<n;++i)
#define reps(i,s,n) for(int i=s;i<n;++i)
using namespace std;
using ll = long long;
using P = pair<int, int>;
using T = tuple<int, int, int>;

int main()
{
   int n;
   cin >> n;
   vector<int> a(n);
   map<int,int> mp;
   rep(i, n)
   {
       cin >> a[i];
       ++mp[a[i]];
   }
   ll ans = 0;
   map<int, int> mp2;

   rep(i, n)
   {
       ++mp2[a[i]];
       ans += n - i - 1 - (mp[a[i]] - mp2[a[i]]);
   }
   cout << ans << endl;

   return 0;
}

あとがき

いつも通りのC問題という感じでしたね。愚直に解けば答えは求まりますが、計算時間が間に合わない、、ということで工夫しましょう！という流れです。B問題にて述べましたがこういうところで計算量という考えがとっても大切になってきます。

でもいきなりは難しいので、まずは

for文が２重になってしまう。２重だとなんかやばそう。

みたいに、ちょっとでも考えてみてください。これを続けていくと大体の見積もりはできるようになってきます。もちろん、制約によっては1重でもダメな時もありますし、３重でも大丈夫な時もあります。このあたりも問題を解いているとだんだんとわかってきます。