백준24261 Same Sum Subsequence

작성일 | In

문제 링크

  • http://icpc.me/24261

시간복잡도

  • $O(N+M)$

풀이

연속한 부분 수열만 고려해도 항상 정답을 찾을 수 있습니다. $A$의 누적 합 배열을 $a$, $B$의 누적 합 배열을 $b$라고 하겠습니다.

일반성을 잃지 않고, $a_n \leq b_m$을 만족한다고 가정합시다. 각 $a_i$에 대해 $b_j \leq a_i$를 만족하는 $j$를 찾으면, $0 \leq a_i - b_j < n$을 만족한다는 것을 알 수 있습니다. $b_j \leq a_i < b_{j+1}$이기 때문입니다.
$a_0-b_0$까지 포함해서 $0 \leq a_i - b_j < n$인 순서쌍 $(i, j)$를 $n+1$개 찾을 수 있습니다. 이때 $a_i - b_j$로 가능한 값의 개수는 $n$가지이므로 비둘기집 원리에 의해 $i’ < i, j’ < j, a_i - b_j = a_{i’} - b_{j’}$을 찾을 수 있습니다. 따라서 항상 연속한 부분 수열 중에 정답이 존재합니다.

전체 코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
#include <bits/stdc++.h>
using namespace std;
using ll = long long;

ll N, M, A[1010101], B[1010101], F;
pair<int,int> Pos[1010101];

void Print(ll *arr, int s, int e){
    cout << e - s + 1 << "\n";
    for(int i=s; i<=e; i++) cout << i - 1 << " ";
    cout << "\n";
}

int main(){
    ios_base::sync_with_stdio(false); cin.tie(nullptr);
    cin >> N;
    for(int i=1; i<=N; i++) cin >> A[i];
    cin >> M;
    for(int i=1; i<=M; i++) cin >> B[i];
    ll SA = accumulate(A+1, A+N+1, 0LL);
    ll SB = accumulate(B+1, B+M+1, 0LL);
    F = SA > SB;
    if(F){
        for(int i=1; i<=max(N,M); i++) swap(A[i], B[i]);
        swap(N, M);
    }

    ll a = 0, b = 0;
    fill(Pos, Pos+1010101, make_pair(-1, 01));
    Pos[0] = {0, 0};
    for(int i=1, j=0; i<=N; i++){
        a += A[i];
        while(j + 1 <= M && b + B[j+1] <= a) b += B[++j];
        int diff = a - b;
        auto pos = Pos[diff];
        if(pos.first == -1){ Pos[diff] = {i, j}; continue; }
        ll s1 = pos.first + 1, e1 = i, s2 = pos.second + 1, e2 = j;
        if(F == 0) Print(A, s1, e1), Print(B, s2, e2);
        else Print(B, s2, e2), Print(A, s1, e1);
        break;
    }
}