查找满足A[i]+A[j]=j-i的数对(i,j)数量

2020年5月3日 malic Comments 0 Comment

问题描述

给定\( N \)个数字组成的数列\( {A_N} \)，\( A_i \)下标从1开始，求出有多少对\( (i,j) \) 使 \( A_i+A_j = j-i \)成立。

https://vjudge.net/problem/AtCoder-5309

例如对于数列 \( \{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)]

Haskell
​x
 
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;
}

C++
 
#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)

Python
 
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

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
​

FAINT

malic自留地

查找满足A[i]+A[j]=j-i的数对(i,j)数量

2020年5月3日 malic Comments 0 Comment

问题描述

问题分析

参考程序

发表回复取消回复

问题描述

问题分析

参考程序

发表回复 取消回复

发表回复取消回复