查找满足A[i]+A[j]=j-i的数对(i,j)数量
问题描述
给定\( N \)个数字组成的数列\( {A_N} \)，\( A_i \)下标从1开始，求出有多少对\( (i,j) \) 使 \( A_i+A_j = j-i \)成立。
例如对于数列 \( \{2, 3, 3, 1, 3, 1\} \)，有三个数对满足条件:
\[A_1+A_4=3=4-1 \\
A_2+A_6=4 = 6-2 \\
A_4+A_6=2=6-4 \]
所以此时结果为3.
问题分析
用枚举当然很简单，可是只有小数据量时能通过。由于此题给出的数列长度达\( 2\times 10^5 \)，\(O(N^2) \)的枚举法会超时。
main = do
line<-getLine
line<-getLine
let p = zip [1..] [read x::Int|x<-words line]
print $ length [1|x<-p,y<-p,fst y > fst x,(fst y)-(fst x)==(snd x)+(snd y)]要想办法优化到最多 \( O(N \log(N) )\)的时间，也就是想办法将逐元素的顺序查找改为二分查找。
我们要对数列\( A_N \)中的两个元素\( A_i,A_j\)符合\[ A_i+A_j = j-i \]
移项即\[ i+A_i = j-A_j \]
令\( L_i = i+A_i \), \(R_i = i-A_i \)，我们只需要对每个\( L \)中的元素\( L_i \)，查找\( R \)中有多少个与 \( L_i \) 相等的即可。估算一下复杂度，对R进行1次排序并进行N次查找， 理论上可以在 \( O(N \log(N) )\)的时间完成。
此题\( A_i \)达\( 10^9 \)，且\( N \) 最大\( 2\times 10^5\)，在计算\( L_i = i+A_i \) 时可能造成int范围溢出，所以将数组设为long long类型
参考程序
#include <cstdio>
#include <algorithm>
using std::sort;
using std::lower_bound;
using std::upper_bound;
const int MAXN=200003;
long long a[MAXN],L[MAXN],R[MAXN];
int main(void)
{
int i,N;
scanf("%d",&N);
for(i=0;i<N;i++)
{
scanf("%lld",a+i);
L[i] = i+1+a[i];
R[i] = (i+1)-a[i];
}
sort(R,R+N);
long long ret = 0,lwbd,upbd;
for(i=0;i<N;i++)
{
lwbd = lower_bound(R,R+N,L[i]) - &R[0];
upbd = upper_bound(R,R+N,L[i]) - &R[0];
ret+=upbd-lwbd;
}
printf("%lld\n",ret);
return 0;
}用Python写时间比较极限，提交了一次运行1996ms，时限就2000ms
def lower_bound(arr,tar):
def _lower_bound(first,length):
if(length==0):
return first
half = length//2
middle = first+half
if arr[middle]<tar:
return _lower_bound(middle+1,length-half-1)
else:
return _lower_bound(first,half)
return _lower_bound(0,len(arr))
def upper_bound(arr,tar):
def _upper_bound(first,length):
if(length==0):
return first
half = length//2
middle = first+half
if arr[middle]>tar:
return _upper_bound(first,half)
else:
return _upper_bound(middle+1,length-half-1)
return _upper_bound(0,len(arr))
N=int(input())
a=list(map(int,input().split()))
L = [(i+1)+a[i] for i in range(N)]
R = [(i+1)-a[i] for i in range(N)]
R.sort()
s=0
for it in L:
s+=upper_bound(R,it)-lower_bound(R,it)
print(s)
Haskell写了一样的方法但是只有小数据能过，还需要再进一步学习
import Data.List
main = do
line<-getLine
line<-getLine
let a = [read x::Int|x<-words line]
let aL = [x+y|(x,y)<-zip a [1..]]
let aR = sort [y-x|(x,y)<-zip a [1..]]
print $ sum $ map (\x ->(upper_bound aR x) - ( lower_bound aR x)) aL
lower_bound :: [Int]->Int->Int
lower_bound a t = _lower_bound a 0 (length a) t
_lower_bound :: [Int]->Int->Int->Int->Int
_lower_bound a first 0 _= first
_lower_bound a first len target = do
let half = div len 2
let middle = first + half
if a!!middle < target
then _lower_bound a (middle+1) (len-half-1) target
else _lower_bound a first half target
upper_bound :: [Int]->Int->Int
upper_bound a t = _upper_bound a 0 (length a) t
_upper_bound :: [Int]->Int->Int->Int->Int
_upper_bound a first 0 _ = first
_upper_bound a first len target = do
let half = div len 2
let middle = first + half
if a!!middle > target
then _upper_bound a first half target
else _upper_bound a (middle+1) (len-half-1) target