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: