ZJNU 2380 - Farmer John Solves 3SUM （区间dp）

USACO 2020 JAN, GOLD - Problem B

ZJNU 2380 / ZJNU contest 1161B

题意

给定长度为N的数组，Q次询问，每次询问给定左右区间 a b

3SUM问题指对于一段区间，选定三个不同位置 i,j,k 使得 a_i+a_j+a_k=0 成立，问这样的三元组组数

对于每次询问输出区间内这样的三元组组数

限制

1≤N≤5000

1≤Q≤10⁵

1≤a_i≤b_i≤N

-10⁶≤A_i≤10⁶

解

仅为解法之一：基于区间长度的区间dp

令 dp[i][j] 表示所求的三元组 (a,b,c) 满足 i≤a,b,c≤j 的组数

区间长度从小到大进行规划，则可知大区间可以由比其小的区间规划而来

所以可以将 dp[i][j] 中的 i 和 j 看作是代表的答案的左右区间

首先得到相邻的 dp[i][j-1]+dp[i+1][j] ，即三元组均在左右区间为 [i,j-1] 以及 [i+1,j] 之内

发现这样计数时，中间的满足左右区间为 [i+1,j-1] 的三元组被重复计数，故需减去dp[i+1][j-1]

最后考虑到这样转移还需要加上表示当前状态的答案，即三元组中有两个固定为 i 和 j 时，满足题意的组数

既然固定了 i 和 j ，表示固定了其中两个数，第三个数 a_k 可以由 a_i+a_j+a_k=0 推得 a_k=-(a_i+a_j)

引入cnt数组动态表示当前 [i+1.j-1] 内各数字出现的次数，则对于状态转移方程，能得到

dp[i][j]=dp[i][j-1]+dp[i+1][j]-dp[i+1][j-1]+cnt[-a[i]-a[j]];

因为数组索引需要为非负数，又考虑到数据范围为 -10⁶≤A_i≤10⁶

所以可以将读入的数全部加上 ave=10⁶ （稍大一点），使其全部成为非负数，再用 cnt[0~2000000] 来表达

每个数都加上基准后，a_i+a_j+a_k=ave*3 ，故第三个数字 a_k=ave*3-(a_i+a_j)

并且注意在使用状态转移方程时判断 a_k 是否会越界（致RE）

---

三元组最小长度为3，故从3开始枚举长度

首先，将 [l+1,r-1] 的数加入cnt数组中

每次枚举，让左边界从1开始，右边界则从len开始，对于cnt即对应 [2,len-1]

每次左边界与右边界需要往右移动一格（窗口整体右移）

所以原本 A[i+1] 的位置会变成左边界，使其从cnt数组中减去

原本 A[j] 的位置会从此时的右边界变成界内元素，故将其加入cnt数组中

整段处理结束后（右边界越界时），需要将cnt数组清零，但直接memset或者遍历清零复杂度很高

所以考虑到最后一次的左右边界分别为 n-len+1 以及 n

所以此时 cnt 数组表示的范围为 [n-len+2,n-1]

又因为转移而使得表示范围变成 [n-len+3,n] ，所以将这一段遍历清零即可

---

处理完dp数组，对于每个询问 l 与 r，输出 dp[l][r] 即可

完整程序

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;

const int ave=1000025,ave3=3000075; //读入数所需要加的基准
int A[5050];
ll dp[5050][5050];
int cnt[2000050];

void solve()
{
    int n,q,l,r,tmp;
    cin>>n>>q;
    for(int i=1;i<=n;i++)
        cin>>A[i],A[i]+=ave;
    for(int len=3;len<=n;len++)
    {
        for(int i=2;i<len;i++)
            cnt[A[i]]++;
        for(int i=1,j=len;j<=n;i++,j++)
        {
            tmp=ave3-(A[i]+A[j]);
            if(tmp>=0&&tmp<2000050)
                dp[i][j]=dp[i][j-1]+dp[i+1][j]-dp[i+1][j-1]+cnt[tmp];
            else
                dp[i][j]=dp[i][j-1]+dp[i+1][j]-dp[i+1][j-1];
            cnt[A[j]]++;
            cnt[A[i+1]]--;
        }
        for(int i=n-len+3;i<=n;i++)
            cnt[A[i]]--;
    }
    while(q--)
    {
        cin>>l>>r;
        cout<<dp[l][r]<<'\n';
    }
}
int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);
    solve();
    return 0;
}