阿里算法岗 0530笔试真题 - 多约束条件下的元素匹配统计

多约束条件下的元素匹配统计

阿里算法岗 0530笔试第二题

题目内容

给定三个长度为 nnn 的数组 {a1,a2,...,an}\{a_1, a_2, \dots, a_n\}{a1,a2,...,an}、{b1,b2,...,bn}\{b_1, b_2, \dots, b_n\}{b1,b2,...,bn} 和 {c1,c2,...,cn}\{c_1, c_2, \dots, c_n\}{c1,c2,...,cn}。这里 cjc_jcj 表示对数组 bbb 的一个下标。请你统计满足以下条件的有序对 (i,j)(i, j)(i,j) 的数量：

1≤i≤j≤n1 \le i \le j \le n1≤i≤j≤n；
ai=bcja_i = b_{c_j}ai=bcj。

输入描述

每个测试文件均包含多组测试数据。第一行输入一个整数 TTT（1≤T≤2×1051 \le T \le 2 \times 10^51≤T≤2×105）代表数据组数；除此之外，保证单个测试文件的 nnn 之和不超过 2×1052 \times 10^52×105。每组测试数据的输入格式如下：

第一行输入一个整数 nnn（1≤n≤1051 \le n \le 10^51≤n≤105）；
第二行输入 nnn 个整数 a1,a2,...,ana_1, a_2, \dots, a_na1,a2,...,an（1≤ai≤1091 \le a_i \le 10^91≤ai≤109）；
第三行输入 nnn 个整数 b1,b2,...,bnb_1, b_2, \dots, b_nb1,b2,...,bn（1≤bi≤1091 \le b_i \le 10^91≤bi≤109）；
第四行输入 nnn 个整数 c1,c2,...,cnc_1, c_2, \dots, c_nc1,c2,...,cn（1≤cj≤n1 \le c_j \le n1≤cj≤n）。

输出描述

对于每一组测试数据，新起一行输出一个整数，表示满足条件的有序对数量。

样例1

输入

复制代码

输出

复制代码

5
6

题解

思路

哈希表

从前往后枚举j,当前可以组成的合法有序对是a下标为(1, j)中值等于b[c[j]]的数量。
按照1的分析，并且随着j递增，不同数字的数量只会增加不会减少，所以可以定义numCount哈希表统计前[1,j]a各种数的数量。
从前往后枚举j，每次j递增时，更新numCount[a[j]]++，并增加当前j情况下能组成的合法有序对ans += numCount[b[c[j]]]

C++

cpp 复制代码

#include<bits/stdc++.h>
using namespace std;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    int T;
    cin >> T;
    while (T--) {
        int n;
        cin >> n;
        vector<int> a(n + 1), b(n + 1), c(n + 1);
        for (int i = 1; i <= n; i++) {
            cin >> a[i];
        }
        for (int i = 1; i <= n; i++) {
            cin >> b[i];
        }
        for (int i = 1; i <= n; i++) {
            cin >> c[i];
        }
        // 记录下标小于j中a每种数字的数量
        unordered_map<int, int> numCount;
        int ans = 0;
        for (int j = 1; j <= n; j++) {
            numCount[a[j]]++;
            ans += numCount[b[c[j]]];
        }
        cout << ans << endl;
    }    
    return 0;
}

java

java 复制代码

import java.util.*;

public class Main {
    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);

        int T = sc.nextInt();

        while (T-- > 0) {
            int n = sc.nextInt();

            int[] a = new int[n + 1];
            int[] b = new int[n + 1];
            int[] c = new int[n + 1];

            for (int i = 1; i <= n; i++) {
                a[i] = sc.nextInt();
            }

            for (int i = 1; i <= n; i++) {
                b[i] = sc.nextInt();
            }

            for (int i = 1; i <= n; i++) {
                c[i] = sc.nextInt();
            }

            // 记录下标小于j中a每种数字的数量
            HashMap<Integer, Integer> numCount = new HashMap<>();

            long ans = 0;

            for (int j = 1; j <= n; j++) {
                numCount.put(a[j], numCount.getOrDefault(a[j], 0) + 1);
                ans += numCount.getOrDefault(b[c[j]], 0);
            }

            System.out.println(ans);
        }

        sc.close();
    }
}

python

python 复制代码

from collections import defaultdict

T = int(input())

for _ in range(T):
    n = int(input())

    a = [0] + list(map(int, input().split()))
    b = [0] + list(map(int, input().split()))
    c = [0] + list(map(int, input().split()))

    # 记录下标小于j中a每种数字的数量
    num_count = defaultdict(int)

    ans = 0

    for j in range(1, n + 1):
        num_count[a[j]] += 1
        ans += num_count[b[c[j]]]

    print(ans)

javascript

js 复制代码

const readline = require('readline');

const rl = readline.createInterface({
    input: process.stdin,
    output: process.stdout
});

const lines = [];

rl.on('line', (line) => {
    lines.push(line);
});

rl.on('close', () => {
    let idx = 0;

    const T = Number(lines[idx++]);

    for (let t = 0; t < T; t++) {
        const n = Number(lines[idx++]);

        const a = [0, ...lines[idx++].split(' ').map(Number)];
        const b = [0, ...lines[idx++].split(' ').map(Number)];
        const c = [0, ...lines[idx++].split(' ').map(Number)];

        // 记录下标小于j中a每种数字的数量
        const numCount = new Map();

        let ans = 0;

        for (let j = 1; j <= n; j++) {
            numCount.set(a[j], (numCount.get(a[j]) || 0) + 1);
            ans += numCount.get(b[c[j]]) || 0;
        }

        console.log(ans);
    }
});

Go

go 复制代码

package main

import "fmt"

func main() {
	var T int
	fmt.Scan(&T)

	for T > 0 {
		T--

		var n int
		fmt.Scan(&n)

		a := make([]int, n+1)
		b := make([]int, n+1)
		c := make([]int, n+1)

		for i := 1; i <= n; i++ {
			fmt.Scan(&a[i])
		}

		for i := 1; i <= n; i++ {
			fmt.Scan(&b[i])
		}

		for i := 1; i <= n; i++ {
			fmt.Scan(&c[i])
		}

		// 记录下标小于j中a每种数字的数量
		numCount := make(map[int]int)

		var ans int64 = 0

		for j := 1; j <= n; j++ {
			numCount[a[j]]++
			ans += int64(numCount[b[c[j]]])
		}

		fmt.Println(ans)
	}
}