ATC-abc417C-Distance Indicators

ATC-abc417C-Distance Indicators

题目大意

给你一个长度为 \(n\) 的整数序列 \(a\)。

找出有多少对整数 \((i,j)\) ，\((1 \le i < j \le n)\) 满足条件 \(j-i=a_i+a_j\) 。

题解

考虑将式子移项，将变量移到两侧。\(a_i+i=a_j-j\) ，由于 \(i<j\) ，所以只需要用树状数组动态维护 \(a_i+i\) 的值，对于 \(j\) ，查询 \(a_j-j\) 即可。

#include<bits/stdc++.h>
#define ios ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);
#define umap unordered_map
#define endl '\n'
using namespace std;
using i128 = __int128;
const int mod =1e9+7;
template <typename T>void read(T&x){
    x=0;int f = 1;
    char c=getchar();
    for(;!isdigit(c);c=getchar())if(c=='-')f=-1;
    for(;isdigit(c);c=getchar())x=(x<<1)+(x<<3)+(c^48);
    x*=f;
}
template <typename T>void print(T x) {
     if (x < 0) { putchar('-'); x = -x; }
     if (x > 9) print(x / 10);
     putchar(x % 10 + '0');
}
#define int long long
class FenwickTree {
private:
    vector<int> tree;  
    int n;  
    int lowbit(int x) {
        return x & -x;
    }

public:
    FenwickTree(int size) : n(size), tree(size + 1, 0) {}

    void add(int pos, int val) {
        for (int i = pos; i <= n; i += lowbit(i)) {
            tree[i] += val;
        }
    }
    int query(int pos) {
        int sum = 0;
        for (int i = pos; i > 0; i -= lowbit(i)) {
            sum += tree[i];
        }
        return sum;
    }
    int query(int left, int right) {
        return query(right) - query(left - 1);
    }
};
const int N=500005;
const int M=2000005;
inline void solve()
{
	int n;
	cin>>n;
	FenwickTree tr(n+5);
	vector<int> num(n+1);
	int ans=0;
	for(int i=1;i<=n;i++) 
	{
		cin>>num[i];
		ans+=tr.query(i-num[i],i-num[i]);
		tr.add(i+num[i],1);
	}
	cout<<ans<<endl;
}

signed main()
{
	ios;
	int T=1;
//	cin>>T;
	for(;T--;) solve();
	return 0;
}