D - Corrupted Array

cpcodeforcesdiv3greedymathconstructive

最終更新日

問題

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

問題文

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

  1. 長さ nn の数列 aa を考える。
  2. bi=ai(1in)b_i = a_i (1 \le i \le n) とする。
  3. bn+1=i=1naib_{n+1} = \sum_{i=1}^n a_i とする。
  4. 1x1091 \le x \le 10^9 となる xx を選び、bn+2=xb_{n+2} = x とする。
  5. bb をシャッフルする。

元の数列 aa としてあり得るのは?

制約

サンプル

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

考察

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

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

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

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

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

コード

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;
}