D - Corrupted Array

問題

https://codeforces.com/contest/1512/problem/D (新しいタブで開く)

問題文

以下のように構築された、長さ $n+2$ の 数列 $b$ が与えられる。

1. 長さ $n$ の数列 $a$ を考える。
2. $b_i = a_i (1 \le i \le n)$ とする。
3. $b_{n+1} = \sum_{i=1}^n a_i$ とする。
4. $1 \le x \le 10^9$ となる $x$ を選び、$b_{n+2} = x$ とする。
5. $b$ をシャッフルする。

元の数列 $a$ としてあり得るのは？

制約

• $1 \le t \le 10^4$
• $1 \le n \le 2 \times 10^5$
• $1 \le b_i \le 10^9$
• $n$ の総和は $2 \times 10^5$ を超えない

サンプル

4
3
2 3 7 12 2
4
9 1 7 1 6 5
5
18 2 2 3 2 9 2
3
2 6 9 2 1

2 3 7
-1
2 2 2 3 9
1 2 6

考察

条件を満たす $a$ が構築できればいいので、要は $\sum_{i=1}^n a_i, x$ を求める問題。
制約から、 $a_i > 0$ が成り立ってるはずなので、 $\forall i (\sum_{j=1}^n a_j > a_i)$ が成り立つ。

まずソートして、 $b'$ とする。
パターンとしては、以下の 2 パターンがあり得る。

• $\sum_{i=1}^n a_i > x$
• $\sum_{i=1}^n a_i \le x$

前者の場合、 $\sum_{i=1}^n a_i = b'_{n+2}$ になる。
あとは、 $x = \sum_{i=1}^{n+1} b'_i - \sum_{i=1}^n a_i$ になるので、これが存在するかチェックすればいい。

後者の場合、 $\sum_{i=1}^n a_i = b'_{n+1} = \sum_{i=1}^n b'_i$ になるので、これが成り立っているかチェックする。

両方成り立たないなら -1 。
あとは $x$ に気を付けて、出力すればいい。

コード

https://codeforces.com/contest/1512/submission/125541260 (新しいタブで開く)

void solve() {
  ll n;
  cin >> n;
  vector<ll> b(n + 2), sum(n + 3, 0);
  rep(i, n + 2) cin >> b[i];
  sort(b.begin(), b.end());
  rep(i, n + 2) sum[i + 1] = b[i] + sum[i];

  ll x = 0;
  ll su = b[n + 1];

  const ll del = sum[n + 1] - b[n + 1];
  rep(i, n + 1) {
    if (b[i] == del) x = b[i];
  }
  if (!x) {
    if (sum[n] == b[n]) {
      su = b[n];
      x = b[n + 1];
    }
    else {
      cout << -1 << endl;
      return;
    }
  }
  bool remX = false;
  bool remS = false;
  rep(i, n + 2) {
    if (!remX && b[i] == x) {
      remX = true;
    }
    else if (!remS && b[i] == su) {
      remS = true;
    }
    else {
      cout << b[i] << " ";
    }
  }
  cout << endl;
}