Educational Codeforces Round 94 (Rated for Div. 2) D题 一维前缀和的另一种用法

https://codeforces.com/contest/1400/problem/D

• 题目很短，意思很清楚，我的想法是记录每一个具有相同元素的区间，记录元素及其位置，然后暴力两两枚举二分位置计算，但这样时间复杂度可能会达到$O(n^3)$，写成下面这样会$TLEontest9$

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;

int main(){
  ios::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);
  int t;
  cin >> t;
  while(t--){
    int n;
    cin >> n;
    map<int, vector<int> > mp;
    vector<int> a(n);
    for(int i=0;i<n;i++){
      cin >> a[i];
      mp[a[i]].push_back(i);
    }
    ll ans = 0;
    for(int i=0;i<n;i++){
      for(int k=i+1;k<n;k++){
        if(a[i] != a[k]) continue;
        for(auto p : mp){
          if(p.first == a[i]){
            int j = upper_bound(p.second.begin(), p.second.end(), i) - p.second.begin();
            int l = upper_bound(p.second.begin(), p.second.end(), k) - p.second.begin();
            int numl = l - 1 - j;
            int numr = (p.second.size() - l);
            if(numl >= 0 && numr >= 0) ans += 1ll * numl * numr;
          }else{
            int j = upper_bound(p.second.begin(), p.second.end(), i) - p.second.begin();
            int l = upper_bound(p.second.begin(), p.second.end(), k) - p.second.begin();
            int numl = l - j;
            int numr = (p.second.size() - l);
            if(numl >= 0 && numr >= 0) ans += 1ll * numl * numr;
          }
        }
      }
    }
    cout << ans << '\n';
  }
  return 0;
}

• 但是考虑到这个题的$a[i]$很小，所以我们可以设$sum[i][j]$表示前$j$个数有多少个$i$，然后枚举每一个可能的位置$j$和$k$，枚举这两个点可以确定剩下的$i$和$l$，如果换一种枚举办法就不太容易，所以我们可以在$O(n^2)$的时间复杂度内得到答案

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;

int main(){
  ios::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);
  int t;
  cin >> t;
  while(t--){
    int n;
    cin >> n;
    vector<int> a(n + 1);
    vector<vector<int> > sum(n + 1, vector<int>(n + 1));
    for(int i=1;i<=n;i++){
      cin >> a[i];
    }
    for(int i=1;i<=n;i++){
      for(int j=1;j<=n;j++){
        sum[i][j] = sum[i][j - 1] + (i == a[j]);
      }
    }
    ll ans = 0;
    for(int j=1;j<=n;j++){
      for(int k=j+1;k<=n;k++){
        ans += 1ll * sum[a[k]][j - 1] * (sum[a[j]][n] - sum[a[j]][k]);
        // a[k]和前j-1个数匹配; a[j]和前k个数匹配
      }
    }
    cout << ans << '\n';
  }
  return 0;
}