[Educatinal Codeforces Round 173(Div.2)] - C. Sums on Segments

https://codeforces.com/contest/2043/problem/C

Educational Codeforces Round 173(Div.2) - C. Sums on Segments

문제 유형: 수학, 부분합, 그리디

시간 제한: 1초

메모리 제한: 256MB

 

문제

You are given an array a of n integers, where all elements except for at most one are equal to −1 or 1. The remaining element x satisfies −10^9≤x≤10^9.

Find all possible sums of subarrays of a, including the empty subarray, whose sum is defined as 0. In other words, find all integers x such that the array a has at least one subarray (possibly empty) with sum equal to x.

A subarray is a contiguous subsegment of an array.

Output these sums in ascending order. Each sum should be printed only once, even if it is achieved by multiple subarrays.

 

n길이의 정수배열을 받습니다. 최대 1개를 제외하고 모두 -1이거나 1입니다. 나머지 요소 x는 -10^9 <= x <= 10^9 입니다.

빈 배열을 포함한 합이 0인 모든 가능한 부분배열의 합을 구하시오. 

x의 모든 가능한 값을 찾아 오른쪽으로 출력하시오.

동일한 합을 여러 부분배열이 가지더라도, 출력은 하나만 되어야합니다.

 

입력

The first line contains a single integer t (1≤t≤104) — the number of test cases.

Then, t test cases follow.

Each test case consists of two lines:

• The first line contains a single integer n(1≤n≤2⋅10^5) — the size of the array.

• The second line contains n integers a1,a2,…,an (−10^9≤ai≤10^9) — the elements of the array a. In the array a, there is at most one element that is neither 1 nor −1.

 

Additional constraint on the input: the sum of n over all test cases does not exceed 2⋅10^5.

 

첫째 줄에 정수 t가 주어집니다. (1 <= t <= 10^4)

t개의 테스트 케이스가 따라옵니다.

각 테스트 케이스는 두 개의 라인이 있습니다.

첫째 줄에 정수 n이 들어옵니다: 배열의 크기

둘때 줄에 각 요소가 n개 있습니다.

 

모든 테스트케이스 n의 합이 2*10^5를 넘지 않ㅎ습니다.

 

출력

For each test case, output two lines:

• In the first line, print a single integer — the number of distinct subarray sums.
• In the second line, print these sums in ascending order.

Each sum should be printed only once, even if it is produced by multiple subarrays.

 

각 테스트 케이스별로, 두 줄을 출력하시오.

- 첫째 줄에 겹치지 않는 부분배열 합들의 개수를 출력하시오.

- 둘째 줄에 각 부분배열 합들을 오름차순으로 출력하시오.

 

풀이

Editorial: https://codeforces.com/blog/entry/137801

x가 0개라면, 반드시 하나의 컴포넌트로 구성된다.

x가 1개라면, 1개 또는 2개의 시퀀스로 구성된다.

 

한 세그먼트의 쌍을 l1, r1, 다른 세그먼트의 쌍을 l2, r2라 한다.

mnl, mxl은 현재까지의 최소/최대 누적합이다.

mnr, mxr은 x이후의 첫 번째 최소/최대 누적합이다.

 

순서대로 배열을 순회하며, [0:i]의 부분합을 구한다.

만약 이번의 수가 x라면, 즉 -1이나 1이 아니라면, 세그먼트의 분리가 필요하다.

mnr에 mnl을, mxr에 mxl을, mnl, mxl에 부분합을 저장한다. 

l1, r1, l2, r2를 각각 적절히 업데이트한다. 자기자신과, 부분합 - (최대l, 최소l, 최대r, 최소r)을 뺀 값과 비교하더 더 작은 값을 넣는다.

그 후, mnl을 더 작은값으로 유지한다.

mxl도 더 큰 값으로 유지한다.

 

최종적으로, l1 ~ r1과 l2 ~ r2의 교집합이 없다면, 각각 배열에 추가하고, 교집합이 있다면, 두 쌍의 최소값부터 최대값까지 넣는다.

 

최대값부터 최소값까지 모두 만들 수 있는 이유는, x 최대 한개를 제외하고는 모두 -1이나 1로 구성되어있기 때문이다. 그래서 최소와 최대값 사이를 모두 만들 수 있다.

 

코드

C++

#include <bits/stdc++.h>

using namespace std;

void solve() {
  int n;
  cin >> n;
  vector<int> arr(n);
  int minSum = 0;
  int maxSum = 0;
  for (int i = 0; i < n; ++i) {
    cin >> arr[i];
  }

  int l1 = 0, r1 = 0;
  int l2 = 2e9, r2 = -2e9;

  int prefixSum = 0;
  int mnl = 0, mxl = 0;
  int mnr = 2e9, mxr = -2e9;

  for (int i = 0; i < n; ++i) {
    prefixSum += arr[i];
    if (arr[i] != -1 && arr[i] != 1) {
      mnr = mnl;
      mxr = mxl;
      mnl = prefixSum;
      mxl = prefixSum;
    }
    l1 = min(l1, prefixSum - mxl);
    r1 = max(r1, prefixSum - mnl);
    l2 = min(l2, prefixSum - mxr);
    r2 = max(r2, prefixSum - mnr);
    mnl = min(mnl, prefixSum);
    mxl = max(mxl, prefixSum);
  }

  vector<int> res;
  if (l2 > r1) {
    for (int i = l1; i <= r1; ++i) {
      res.push_back(i);
    }
    for (int i = l2; i <= r2; ++i) {
      res.push_back(i);
    }
  } else if (r2 < l1) {
    for (int i = l2; i <= r2; ++i) {
      res.push_back(i);
    }
    for (int i = l1; i <= r1; ++i) {
      res.push_back(i);
    }
  } else {
    for (int i = min(l1, l2); i <= max(r1, r2); ++i) {
      res.push_back(i);
    }
  }

  cout << res.size() << "\n";
  for (int i = 0; i < res.size(); ++i) {
    cout << res[i] << " ";
  }
  cout << "\n";
}

int main(int argc, char *argv[]) {
  int t;
  cin >> t;
  while (t--) {
    solve();
  }
  return 0;
}