Official

C - Made Up Editorial by en_translator


We want to count the number of \(j\) such that \(A_i = B_{C_j}\) for a fixed \(i\). How can we precalculate this?

Since the value of \(B_k\) is between \(1\) and \(N\), we can prepare an array of length about \(N\), and try every value of \(j\), so that we can store the number of \(j\) such that \(\mathrm{count}_x\) for each \(1 \leq x \leq N\) as an array \(\mathrm{count}_x\).

The answer is the sum of \(\mathrm{count}_{A_i}\) for each \(i\). Note that the answer may now fit in a \(32\)-bit integer, since the answer can be up to \(N^2\).

As you can see, for a problem like “repeating the same process many times” can often be handled with precalculations.

Sample Code in 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';
}

Sample Code in 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: