合計が0になるように並べ替え可能か判定するC++コード

問題の概要

-1 から 1 の範囲の値を持つ n 個の要素からなる配列 A と、Q[i] = (li, ri) の形式で表される m 個のクエリ Q が与えられます。各クエリに対しては、配列の要素を並べ替えることで区間 Q[li] ～ Q[ri] の合計が 0 になるようにできる場合は 1、できない場合は 0 を答えとします。すべてのクエリに対する回答を求めましょう。

たとえば、入力が A = [-1, 1, 1, 1, -1]、Q = [[1, 1], [2, 3], [3, 5], [2, 5], [1, 5]] のとき、出力は [0, 1, 0, 1, 0] となります。

解法の考え方

この問題は次の手順で解くことができます。

1. 配列内に含まれる -1 の個数を数えて z とします。
2. -1 と 1 をペアにできる最大数は、両者の個数のうち小さい方です。そこで z > n - z の場合は、z を n - z に置き換えます。
3. 各クエリ (l, r) について、次の 2 つの条件を満たせば 1 を、そうでなければ 0 を出力します。
   • 区間の長さ (r - l + 1) が偶数であること
   • 区間の長さの半分 (r - l + 1) / 2 が z 以下であること

アルゴリズムの手順

n := size of A
m := size of Q
z := 0
for initialize i := 0, when i < n, update (increase i by 1), do:
    z := z + (1 if A[i] < 0, otherwise 0)
if z > n - z, then:
    z := n - z
for initialize i := 0, when i < m, update (increase i by 1), do:
    l := Q[i, 0]
    r := Q[i, 1]
    print 1 if ((r - l + 1) mod 2 is 0 and (r - l + 1) / 2 <= z),
otherwise 0

C++実装例

理解を深めるために、実際の実装を見てみましょう。

#include <bits/stdc++.h>
using namespace std;

void solve(vector<int> A, vector<vector<int>> Q){
    int n = A.size();
    int m = Q.size();
    // 配列内の -1 の個数を数える
    int z = 0;
    for (int i = 0; i < n; ++i)
        z += A[i] < 0;
    // 作れるペアの最大数は -1 と 1 の個数の少ない方
    if (z > n - z)
        z = n - z;
    for (int i = 0; i < m; i++){
        int l = Q[i][0];
        int r = Q[i][1];
        // 長さが偶数 かつ 必要なペア数が z 以下なら 1
        bool ok = ((r - l + 1) % 2 == 0) && ((r - l + 1) / 2 <= z);
        cout << ok << ", ";
    }
}

int main(){
    vector<int> A = { -1, 1, 1, 1, -1 };
    vector<vector<int>> Q = { { 1, 1 }, { 2, 3 }, { 3, 5 }, { 2, 5 }, { 1, 5 } };
    solve(A, Q);
}

入力

{ -1, 1, 1, 1, -1 }, { { 1, 1 }, { 2, 3 }, { 3, 5 }, { 2, 5 }, { 1, 5 } }

出力

0, 1, 0, 1, 0,

コードのポイント

区間の合計を 0 にするためには、区間内の -1 と 1 の個数がちょうど等しくなる必要があります。そのため、区間の長さは必ず偶数でなければなりません。さらに、必要なペアの数（区間の長さの半分）が、配列全体で作れる -1 と 1 のペアの最大数 z を超えていてはなりません。

この方法なら、前処理に O(n)、各クエリの判定に O(1) しかかからないため、クエリが多数あっても効率的に処理できます。全体の計算量は O(n + m) です。