合法元组数(C++ 记录先前遍历过的数据)

题目描述
给一个长度为 n 的序列a[n], 请问有多少个三元组(i, j, k)满足i < j < k 且a[i] - a[j] = 2a[j] - a[k]？输出符合的数量。
 
 
输入描述

第一行一个整数 n，表示序列长度;
第二行用空格隔开n个整数，表示a[i];
1 <= n <= 4000, 0 <= a[i] <= 1,000,000

输出描述

一行一个整数，表示符合的三元组的数量

样例输入

4
4 2 2 2

样例输出

3

🙏分析🙏

可以将之前遍历过的数字的个数用map存起来，因为题目只要求个数，不要求具体的组合，从而减少了遍历数，复杂度变成$O(n^2)$。

代码

#include <iostream>
#include <unordered_map>

using namespace std;

int main() {
    int arr[40000] = {0};
    unordered_map<int, int> mp;
    long long ans = 0;
    int n;
    cin >> n;
    if (n < 3) {
        cout << 0;
        return 0;
    }
    for (int i = 0; i < n; i++) {
        cin >> arr[i];
    }
    
    for(int j = 0; j < n; j++){ // 从 j = 0 开始, 不是 1
    							// 要记录mp[arr[0]]
    							// 或者在for循环之前就给mp[arr[0]]赋值也行
        for(int k = j + 1; k < n; k++){
            int t = 3 * arr[j] - arr[k];
            if(mp.count(t)) {
                ans += mp[t];
            }
        }
        mp[arr[j]]++; // 这行是关键, 看清楚是[j],
        // “之前”，指的是当前的 j和k 之前
        // 有因为j < k, 那就是 j 之前
    }
    cout << ans;
    return 0;
}