Pythonでボールが入る箱の行と列の位置を効率的に求める方法

2つの配列AとBがあるとします。配列Aのサイズは行数を表し、A[i]はi行目に存在する箱の個数を意味します。一方、配列Bはボールの配列であり、B[i]はボールに書かれた番号を表します。各ボールi（値がB[i]）は、先頭から数えてB[i]番目の箱に配置されるものとします。このとき、B[i]のそれぞれに対応する箱の「行」と「列」を求めるのが本記事の目的です。

例えば、入力が A = [3, 4, 5, 6]、B = [1, 3, 5, 2] の場合、出力は [(1, 1), (1, 3), (2, 2), (1, 2)] となります。具体的な対応関係は以下の通りです。

• B[0] = 1 → 1行目・1列目
• B[1] = 3 → 1行目・3列目
• B[2] = 5 → 2行目・2列目
• B[3] = 2 → 1行目・2列目

解決のアプローチ：累積和と二分探索

この問題を効率的に解く鍵となるのは、累積和（prefix sum）と二分探索（binary search）の組み合わせです。配列Aを累積和に変換しておけば、「ある番号のボールが何行目に入るか」を高速に特定できます。手順は以下の通りです。

1. len_a を配列Aのサイズ、len_b を配列Bのサイズとします。
2. i を 1 から len_a までループし、A[i] = A[i] + A[i-1] として累積和を作成します。
3. i を 0 から len_b までループし、以下の処理を行います。
   • row := ソート順を保ったままB[i]を挿入できるインデックス（二分探索で取得）
   • row >= 1 の場合：box_num := B[i] − A[row − 1]
   • それ以外の場合：box_num := B[i]
   • (row + 1, box_num) のペアを出力します。

Pythonでの実装例

Pythonでは標準ライブラリの bisect モジュールを使うことで、二分探索を簡単に実装できます。以下がコード例です。

import bisect

def get_position(A, B):
    len_a = len(A)
    len_b = len(B)
    # 配列Aを累積和に変換
    for i in range(1, len_a):
        A[i] += A[i - 1]
    # 各ボールの行と列を求める
    for i in range(len_b):
        row = bisect.bisect_left(A, B[i])
        if row >= 1:
            box_num = B[i] - A[row - 1]
        else:
            box_num = B[i]
        print((row + 1, box_num))

A = [3, 4, 5, 6]
B = [1, 3, 5, 2]
get_position(A, B)

入力

[3, 4, 5, 6], [1, 3, 5, 2]

出力

(1, 1)
(1, 3)
(2, 2)
(1, 2)

処理のポイント

累積和変換後の配列Aは [3, 7, 12, 18] となり、各行までの箱の総数を表します。bisect_left により、ボールの番号がどの区間に属するかをO(log n)で判定でき、その区間の開始値との差分が列番号になります。この手法により、ボールごとに全行を走査するO(n×m)の素朴な方法と比べ、大幅な高速化が可能です。