【C++】2つの配列を等しくするために必要な操作手順を求めるアルゴリズム

問題概要

n個の要素を持つ2つの配列 A と B があるとします。ここで、次のような操作を考えます。

操作： 2つのインデックス i と j を選び、i 番目の要素を1減らし、j 番目の要素を1増やします。ただし、操作を行った後も配列のすべての要素は非負（0以上）でなければなりません。

この操作を繰り返して配列 A と B を同じ内容にしたい場合、必要な操作の手順（インデックスのペア）を求めてください。どうしても一致させられない場合は -1 を返します。

たとえば、入力が A = [1, 2, 3, 4]、B = [3, 1, 2, 4] のとき、出力は [(1, 0), (2, 0)] となります。これは、i = 1、j = 0 の操作によって配列が [2, 1, 3, 4] となり、続いて i = 2、j = 0 の操作によって [3, 1, 2, 4] となるためです。

解法のアプローチ

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

1. 総和の確認： まず、配列 A と B の要素の総和をそれぞれ計算します。総和が異なる場合、どれだけ操作を繰り返しても両者を一致させることはできません（操作しても総和は変化しないため）。この場合は -1 を返します。
2. 差分の記録： 各インデックスについて差分 C[i] = A[i] - B[i] を計算して配列 C に保存し、同時に絶対値の合計を求めます。この合計を2で割った値が、必要な操作回数になります。
3. 操作の出力： 差分が正のインデックス（要素を減らす側）と負のインデックス（要素を増やす側）を順番に探索しながら、操作のペアを出力します。

アルゴリズムの擬似コード

a := 0, b := 0, c := 0
n := size of A
Define an array C of size n and fill with 0
for initialize i := 0, when i < n, update (increase i by 1), do:
    a := a + A[i]
for initialize i := 0, when i < n, update (increase i by 1), do:
    b := b + B[i]
if a is not equal to b, then:
    return -1
Otherwise
    for initialize i := 0, when i < n, update (increase i by 1), do:
        c := c + |A[i] - B[i]|
        C[i] := A[i] - B[i]
    c := c / 2
    i := 0
    j := 0
    while c is non-zero, decrease c after each iteration, do:
        while C[i] <= 0, do:
            (increase i by 1)
        while C[j] >= 0, do:
            (increase j by 1)
        print i and j
        decrease C[i] and increase C[j] by 1

C++による実装例

それでは、実際のC++コードを見てみましょう。

#include <bits/stdc++.h>
using namespace std;
void solve(vector<int> A, vector<int> B){
    int a = 0, b = 0, c = 0;
    int n = A.size();
    vector<int> C(n, 0);
    // 配列Aの総和を計算
    for (int i = 0; i < n; i++)
        a += A[i];
    // 配列Bの総和を計算
    for (int i = 0; i < n; i++)
        b += B[i];
    // 総和が異なる場合は一致させられない
    if (a != b){
        cout << -1;
        return;
    }
    else{
        // 差分の絶対値の合計と差分配列を計算
        for (int i = 0; i < n; i++){
            c += abs(A[i] - B[i]);
            C[i] = A[i] - B[i];
        }
        c = c / 2; // 必要な操作回数
        int i = 0, j = 0;
        while (c--){
            // 減らす対象（差分が正）を探す
            while (C[i] <= 0)
                i++;
            // 増やす対象（差分が負）を探す
            while (C[j] >= 0)
                j++;
            cout << "(" << i << ", " << j << "), ";
            C[i]--, C[j]++;
        }
    }
}
int main(){
    vector<int> A = { 1, 2, 3, 4 };
    vector<int> B = { 3, 1, 2, 4 };
    solve(A, B);
}

実行結果

入力

{ 1, 2, 3, 4 }, { 3, 1, 2, 4 }

出力

(1, 0), (2, 0),

計算量について

このアルゴリズムの時間計算量は O(n + c) です。ここで n は配列のサイズ、c は必要な操作回数（差分の絶対値の合計の半分）を表します。空間計算量は差分配列 C を格納するために O(n) となります。

なお、元のコードには配列 B の総和を計算する箇所で誤って A[i] を加算するバグがありましたが、上記の実装では正しく B[i] を加算するよう修正しています。この修正がないと、総和の比較が常に成立してしまい、本来 -1 を返すべきケースでも誤った動作をする可能性があります。