C++で同じ配列内に平均値が存在するペアを数える方法

問題の概要

各要素が[-1000, 1000]の範囲に収まる整数型配列が与えられます。求めるのは、「2つの要素の平均値が同じ配列内にも存在する」ようなペアの総数です。

例えば、arr[] = [1, 2, 3, 4] の場合を考えてみましょう。1と3の平均は2、2と4の平均は3であり、2も3も配列内に存在するため、該当するペアは (1, 3) と (2, 4) の2つとなり、答えは 2 になります。

入出力例

入力: arr[] = { -1, 2, 5, -3, 8, 10 }
出力: 同じ配列内に平均値が存在するペアの数: 2
説明: 平均値がarr[]内に存在するペアは (-1, 5)(平均は2)と (2, 8)(平均は5)の2つです。

入力: arr[] = { 1, 3, 2, 5, 10, 6 }
出力: 同じ配列内に平均値が存在するペアの数: 3
説明: 該当するペアは (1, 3)(平均は2)、(1, 5)(平均は3)、(2, 10)(平均は6)の3つです。

アルゴリズムの考え方

このプログラムでは頻度配列(度数分布)を活用します。負の要素も扱う必要があるため、頻度配列のサイズは「範囲の2倍+1」である2001とします。

• 負の数の頻度はインデックス0〜1000に格納
• 正の数の頻度はインデックス1000〜2000に格納

具体的には、元の値にN(=1000)を加算した位置をインデックスとして使用することで、負の値も正の値も同じ配列で管理できます。

非ゼロの頻度に対する処理

1. 同一値同士のペア: countに freq[i] × (freq[i] - 1) / 2 を加算します。同じ数同士の平均はその数自身になるためです。例えば配列に2が5個ある場合、ペアの総数は 5 × (5 - 1) / 2 = 10 となります。
2. 1つ飛ばしの走査: freq[i]が非ゼロの場合、隣接する数ではなく1つ飛ばしで頻度を調べます。連続する2数の平均は小数になり、整数配列内には存在しないためです。
3. 平均の存在確認: freq[j]が非ゼロで、かつ freq[(i + j) / 2](2数の平均に対応する頻度)も非ゼロであれば、countに freq[i] × freq[j] を加算します。各数は相手側のすべての数とペアになり得るためです。

手順の詳細

1. 整数型配列 arr[] を受け取ります。
2. 関数 average_pair(arr, size) は、ペアの平均値が配列arr[]内にも存在するようなペアの数を返します。
3. 初期カウントを0とし、N = 1000 で初期化します。
4. 頻度配列の長さを size_2 = 2 × N + 1(範囲[-1000, 1000]用)として計算します。
5. 頻度配列をすべて0で初期化します。
6. 負の要素の頻度はインデックス0〜1000に、正の要素はそれ以降に格納されるよう、インデックスにNを加算して頻度配列を作成します。
7. 各要素 arr[i] について、arr_freq[arr[i] + N]++ で頻度を更新します。
8. i = 0 から i < size_2 まで頻度配列を走査します。
9. 非ゼロの頻度ごとに、条件1に従って freq[i] × (freq[i] - 1) / 2 をcountに加算します。
10. arr_freq[i]が非ゼロなので、1つ飛ばしで他の頻度を走査します。
11. temp_2 を arr_freq[(i + j) / 2] として計算します。
12. temp_2が非ゼロかつ arr_freq[j]が非ゼロなら条件3を満たすので、countに積 arr_freq[i] × arr_freq[j] を加算します。
13. すべての反復が終わると、countに該当ペアの総数が格納されます。
14. 結果としてcountを返します。

C++での実装例

#include <bits/stdc++.h>
using namespace std;
int average_pair(int arr[], int size_1){
    int count = 0;
    int N = 1000;
    int size_2 = (2 * N) + 1;
    int arr_freq[size_2] = { 0 };
    for (int i = 0; i < size_1; i++){
        int temp = arr[i];
        arr_freq[temp + N]++;
    }
    for (int i = 0; i < size_2; i++){
        if (arr_freq[i] > 0){
            int check = (arr_freq[i]) * (arr_freq[i] - 1);
            count += check / 2;
            for (int j = i + 2; j < 2001; j += 2){
                int temp_2 = arr_freq[(i + j) / 2];
                if (arr_freq[j] > 0 && temp_2 > 0){
                    count += (arr_freq[i] * arr_freq[j]);
                }
            }
        }
    }
    return count;
}
int main(){
    int arr[] = { 2, 3, 1, 8, 9, 10 };
    int size = sizeof(arr) / sizeof(arr[0]);
    cout<<"Count of pairs with average present in the same array are: "<<average_pair(arr, size);
    return 0;
}

出力結果

上記のコードを実行すると、次の出力が得られます。

Count of pairs with average present in the same array are: 2

計算量について

頻度配列の構築には配列の要素数nに対してO(n)、その後の走査は頻度配列のサイズM(=2001)に対して各インデックスから最大約M/2回の比較を行うため、全体の時間計算量はO(M²)、空間計算量はO(M)となります。要素の値域が限定されている場合、要素数が多くても安定した性能を発揮できるのがこの手法の大きな利点です。