ABC202 C 解答

C - Made Up（204）

問題

C - Made Up

問題文

1 以上 N 以下の整数からなる長さ N の数列 
A = ( A1, A2, … , AN ), B = ( B1, B2, … , BN ), C = ( C1, C2, … , CN ) が与えられます。
1 以上 N 以下の整数 i, j の組 ( i, j ) であって、 Ai = B_{Cj} となるものの総数を求めてください。

制約

1 ≤ N ≤ 10^5
1 ≤ Ai, Bi, Ci ≤ N
入力は全て整数である。

考察

答えを求めるだけなら簡単ですね。全部のインデックスについて調べ上げて、一致していたら答えを+1してあげればいずれは答えが求まります。しかし、この方法ですとfor文を2重で回さなければならないので計算量はO(N)です。制約をみるに間に合わなさそうです。

ですので、着目点を変えてみます。今回は

どの数字でAとBが一致するか

ということを考えていきます。

A : 1 2 2
B : 3 1 2
C : 2 3 2

どの場所で一致するか(for文が2重になる)

どの数字が一致するか（今回考える）

（Cを使って、Bの配列を新しく作っています。Cの数字の場所のBの数字を取り出せば作れます。）

下の図を見ると、１が2回、２が２回一致していますね。ということで、答えは４回です。この例は入力例１ですので、AtCoder様のページを見に行くとどうやらこれであっていそうです。

さて、もう少し下の例で考えてみます。

まずは、１が一致する組み合わせがいくつあるかを求めます。下の例ではAの配列（上）では１が1回だけ出現しています。一方で、Bの配列（下）では２回出現していますね。上の１に対して、下は2つの１が対応しますので、その総数は

1*2 = 2通り

ですね。２も同様に考えると

2 * 1 = 2通り

となります。各数字において

Aの出現数 * Bの出現数

がその数字にて一致するi, j の組み合わせの数になります。

各数字におけるAとBの一致する通り数を求めてあげたらあとはそれを全ての数字に対して足し合わせることで答えが求まります。

ですので、答えを求めるために必要なのは

Cを使ってBの配列を新しくする

AとBにおいて各数字が何回出現したかを求める

です。この２つでしたら、どちらもfor文の1重ループで求めることができます。ということで、時間内で問題を解くことができました。 

実装

 #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<ll, ll>;

const int N = 1e5;
int main()
{
   int n;
   cin >> n;
   vector<int> a(n), b(n), c(n);
   rep(i, n) cin >> a[i];
   rep(i, n) cin >> b[i];
   rep(i, n) cin >> c[i];

   vector<int> A(N + 5), B(N + 5);
   rep(i, n) ++A[a[i]];
   rep(i, n) ++B[b[c[i]-1]];
   ll ans = 0;
   rep(i, N+1) ans += (ll)A[i] * B[i];
   cout << ans << endl;
   return 0;
}

あとがき

A, B に出現する各数字を求めるときには制約に注目してあげましょう。今回各数字は

1 ≤ N ≤ 10^5
1 ≤ Ai, Bi, Ci ≤ N

まで出現するみたいです。ですので、実装では一番上に宣言しています。ただ、お恥ずかしながら私はコンテスト中にこの数字を間違えてWAを出しました。

数字の上限を気にしない方法としてmapを使うというものがあります。若干計算量は多くなりますが、十分に間に合う程度です。ミスをするぐらいでしたら、便利なデータ構造を使いましょう。と書いていますが、これは自分に強く言い聞かせています。