Official
C - Made Up Editorial by KoD
\(i\) を固定したときに、\(A_i = B_{C_j}\) となる \(j\) の個数が高速に求められればよいです。これを前もって計算しておく方法を考えます。
\(B_k\) の値は \(1\) 以上 \(N\) 以下であるので、長さ \(N\) 程度の配列 \(\mathrm{count}_x\) を用意し、\(j\) の値を全て試すことによって、各 \(1 \leq x \leq N\) について \(B_{C_j} = x\) となる \(j\) の個数を \(\mathrm{count}_x\) として保持しておくことが出来ます。
各 \(i\) について \(\mathrm{count}_{A_i}\) を足し合わせたものが答えとなります。答えは最大で \(N^2\) になるため、\(32\) bit 整数に収まらないことに注意してください。
このように、「同じような処理を何度も行う」ような問題においては、「前もって計算しておく」という方法がうまくいくことがしばしばあります。
C++ での実装例:
#include <iostream>
#include <vector>
int main() {
int N;
std::cin >> N;
std::vector<int> A(N), B(N), C(N);
for (auto& x : A) {
std::cin >> x;
x -= 1;
}
for (auto& x : B) {
std::cin >> x;
x -= 1;
}
for (auto& x : C) {
std::cin >> x;
x -= 1;
}
std::vector<int> count(N);
for (int j = 0; j < N; ++j) {
count[B[C[j]]] += 1;
}
long long ans = 0;
for (int i = 0; i < N; ++i) {
ans += count[A[i]];
}
std::cout << ans << '\n';
}
Rust での実装例:
use proconio::{input, marker::Usize1};
fn main() {
input! {
n: usize,
a: [Usize1; n],
b: [Usize1; n],
c: [Usize1; n],
}
let mut count = vec![0; n];
for x in c {
count[b[x]] += 1;
}
println! {
"{}",
a.into_iter()
.map(|x| count[x])
.sum::<u64>()
}
}
posted:
last update: