Codeforces Round #707 C. Going Home

C. Going Home

题意

给 $n(1\le n\le 2\cdot 10^5)$ 个数，每个数 $1\le a_i\le 2.5\cdot 10^6$ ，要求找到四个不同的数 $x,y,z,w$ ，使得 $a_x+a_y=a_z+a_w$ 。不能找到则输出 NO。

题解

• 注意到 $a_i$ 的范围为 $1\le a_i\le 2.5\cdot 10^6$ ，可以先找两个数的和为某个值的所有组合，即 $v[x]$ 存储所有和为 $x$ 的一对下标，但是要求四个下标都不同，所以可能找到很多无效的下标对。这样的复杂度为 $O(C\cdot n^2)$，$C$ 表示有相同下标的下标对个数。
• 为了减少无效的下标对个数，可以先做预处理。记录每一个值出现的所有位置，如果一个值出现了 $4$ 次，那么这四个位置就是答案。如果有两个值都出现了两次，那么这四个位置也可以是答案。那么剩下的情况最多有一对数出现超过两次，其他的数各不相同。也就是最多找到一次无效的答案。这样就复杂度变为 $O(n^2)$ 。
• 两个数的和的值域为 $[2,5e6]$，那么找大约 $5e6$ 次一定可以找到答案，复杂度为 $O(min(n^2,5e6))$。

代码

#pragma region
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <vector>
using namespace std;
typedef long long ll;
#define rep(i, a, n) for (int i = a; i <= n; ++i)
#define per(i, a, n) for (int i = n; i >= a; --i)
#pragma endregion
const int maxn = 2e5 + 5;
vector<pair<int, int>> v[int(5e6) + 5];
vector<int> pos[int(5e6) + 5];
int a[maxn];
int main() {
    
    
    int n;
    scanf("%d", &n);
    int c = 0;
    rep(i, 1, n) {
    
    
        scanf("%d", &a[i]);
        pos[a[i]].push_back(i);
        if (pos[a[i]].size() == 4) {
    
    
            puts("YES");
            for (auto it : pos[a[i]]) printf("%d ", it);
            return 0;
        }
        if (pos[a[i]].size() == 2) {
    
    
            if (c) {
    
    
                puts("YES");
                printf("%d %d %d %d\n", pos[c][0], pos[a[i]][0], pos[c][1], pos[a[i]][1]);
                return 0;
            }
            c = a[i];
        }
    }
    rep(i, 1, n) rep(j, i + 1, n) {
    
    
        for (auto it : v[a[i] + a[j]]) {
    
    
            set<int> st;
            st.insert(i), st.insert(j), st.insert(it.first), st.insert(it.second);
            if (st.size() == 4) {
    
    
                printf("YES\n%d %d %d %d\n", i, j, it.first, it.second);
                return 0;
            }
        }
        v[a[i] + a[j]].push_back({
    
    i, j});
    }
    puts("NO");
}