Pythonで隠し配列から最頻出要素のインデックスを求めるプログラムの実装方法

問題概要

ここでは、「TestArray」というクラスが与えられた状況を考えます。このクラスは、値として 0 か 1 のみを格納できる非公開（private）の配列を内部に持ち、外部から利用できる公開メンバー関数として length() と query() の2つを提供しています。

• length()：配列の長さを返します。

• query(p, q, r, s)：4つのインデックスを受け取り、それぞれの位置にある値を比較して、次の3種類の値のいずれかを返します。

  • 指定された4つのインデックスの値がすべて同じ（すべて 0、またはすべて 1）場合 → 4 を返す

  • 3つの値が同じで、残りの1つだけが異なる場合 → 2 を返す

  • 0 が2つ、1 が2つ含まれる場合 → 0 を返す

やりたいこと

配列本体には一切直接アクセスせず、クラスが公開しているメンバー関数だけを使って、配列内で最も多く出現する要素（0 または 1）のインデックスを特定します。0 と 1 の出現回数が同じ場合は -1 を返します。

たとえば、入力が array = [0, 1, 1, 0, 1, 1, 1, 0] の場合、出力は 2 になります。インデックス 2 の値は 1 であり、1 がこの配列で最も頻出する値だからです。同様に、1・4・5・6 も正解として成立します（これらのインデックスにも値 1 が格納されているためです）。

解法のステップ

この問題は、以下の手順に従って解きます。

1. n := length() で配列の長さを取得します。

2. groupA := 1、groupB := 0、aIdx := null、bIdx := null で変数を初期化します。

3. first := query(0, 1, 2, 3)、second := query(0, 1, 2, 4) をあらかじめ計算しておきます。

4. i を 4 から n-1 までループします。query(0, 1, 2, i) の結果が first と同じなら groupA を +1 して aIdx := i、異なれば groupB を +1 して bIdx := i とします。

5. i を 0 から 2 までループします。[0, 1, 2, 3, 4] から i を除いた4つのインデックスで query を実行し、その結果が second と同じなら groupA を +1 して aIdx := i、異なれば groupB を +1 して bIdx := i とします。

6. 最後に、groupA > groupB なら aIdx を、groupB > groupA なら bIdx を返します。両者が等しい場合は -1 を返します。

このアプローチのポイントは、各要素を「query の結果が基準と一致するグループ」「一致しないグループ」の2つに振り分け、サイズが大きい方のグループに属するインデックスを答えとしている点です。0 と 1 の個数がちょうど半々のときにのみ -1 が返る仕組みになっています。クエリ回数は要素数に対して線形（O(n)）なので、非常に効率的な解法といえます。

Pythonでの実装例

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

class TestArray:
    def __init__(self, array) -> None:
        self.__arr = array

    def length(self):
        return len(self.__arr)

    def query(self, p, q, r, s):
        val = self.__arr[p] + self.__arr[q] + self.__arr[r] + self.__arr[s]
        if val == 4 or val == 0:
            return 4
        elif val == 1 or val == 3:
            return 2
        elif val == 2:
            return 0

def solve(reader):
    n,groupA,groupB,aIdx,bIdx=reader.length(),1,0,None,None
    first,second=reader.query(0,1,2,3),reader.query(0,1,2,4)
    for i in range(4,n):
        if reader.query(0,1,2,i)==first:
            groupA,aIdx=groupA+1,i
        else:
            groupB,bIdx=groupB+1,i
    for i in range(3):
        nxt=[v for v in [0,1,2,3,4] if v!=i]
        if reader.query(*nxt)==second:
            groupA,aIdx=groupA+1,i
        else:
            groupB,bIdx=groupB+1,i
    return aIdx if groupA>groupB else bIdx if groupB>groupA else -1

arr_ob = TestArray([0, 1, 1, 0, 1, 1, 1, 0])
print(solve(arr_ob))

入力

[0, 1, 1, 0, 1, 1, 1, 0]

出力

2