python 配列 高速ソート

Question

次のような配列xがあり(実際に使用するデータはもっと大きい)，そのデータから次のソースコードで得られる結果resul1,result2を得たいです．実際のデータにおいても配列x[0],x[1]...はそれぞれソートされているものとします．
得たい結果は

• xの全要素をソートした結果のうち5以下のものからなる配列result1
• result1の要素がxの何番目の要素であったかを表す配列result2

です．

次のソースコードを用いても欲しい結果は得られますが，扱うデータが巨大なため，実行に時間がかかりすぎてしまいます．
リスト内包表記を用いる方法など，何か高速な方法はありますでしょうか？

扱うデータ

import numpy as np

a = np.array([1, 2.3, 3.1, 4.6, 6.7])
b = np.array([2.3, 4.2])
c = np.array([])
d = np.array([2.5, 3.2, 4.3, 4.4, 4.8, 5.0])
x = [a, b, c, d]

自作コード

import numpy as np

result1 = []
result2 = []
append1 = result1.append
append2 = result2.append

while True:
  value = 6
  for i in range(len(x)):
    if x[i].size and value > x[i][0]:
      value = x[i][0]
      index = i

  if value > 5:
    break

  append1(value)
  append2(index)
  x[index] = np.delete(x[index], 0)

自作コードで得られる結果かつ得たい結果
result1=[1.0, 2.3, 2.3, 2.5, 3.1, 3.2, 4.2, 4.3, 4.4, 4.6, 4.8, 5.0]
result2=[0, 0, 1, 3, 0, 3, 1, 3, 3, 0, 3, 3]

以前にも似たような質問をさせていただきましたが，xのそれぞれの要素x[0],x[1]...がソート済みであると明示せずに，質問させていただいたので，ソート済みであればより早い方法があるのではないかと思い質問させていただきました．

Answer 1

データ量が10**6程度までであれば、過去のやりとりで出てきているnumpy.argsortで十分高速だと思います。

import numpy as np

# オリジナル
def solve1(x):
    result1 = []
    result2 = []
    append1 = result1.append
    append2 = result2.append

    while True:
        value = 6
        for i in range(len(x)):
            if x[i].size and value > x[i][0]:
                value = x[i][0]
                index = i
        if value > 5:
            break

        append1(value)
        append2(index)
        x[index] = np.delete(x[index], 0)
    return result1, result2

# numpy.argsort()使用
def solve2(x, limit=5.0):
    x = [arr[:np.searchsorted(arr, limit, side='right')] for arr in x]
    result1 = np.concatenate(x)
    result2 = np.concatenate([np.full(arr.size, i) for i, arr in enumerate(x)])
    idx = np.argsort(result1, kind='mergesort')
    return result1[idx], result2[idx]


np.random.seed(123)
size = 10**5  # テスト用のデータサイズ
a = np.sort(np.random.rand(size) * 10)
b = np.sort(np.random.rand(size) * 10)
c = np.sort(np.random.rand(size) * 10)
d = np.sort(np.random.rand(size) * 10)

x = [a, b, c, d]
r11, r12 = solve1(x)

x = [a, b, c, d]
r21, r22 = solve2(x)

print(r11 == r21.tolist() and r12 == r22.tolist()) # 結果が同じかチェック

Comments

[bobslay]

ご回答ありがとうございます．
実際に使用したいデータを見たところ，サイズが10^10ほどでした．
そもそもデータ量が大きすぎてこれ以上高速に実行できないのではないかと思ってきました．．．

[退会済みユーザー]

10^10は大きいですね。確認ですが、お使いのマシンでは10^10のデータでもオンメモリで処理できるだけのRAMを搭載しているんですよね? 私の環境(RAM 32GB)だと size=10**8で処理時間が4.5秒、size=5*10**8で144秒(ただしRAM不足でスワップ発生)でした。データサイズを少しずつ大きくしながら試してみてはどうでしょうか。

Answer 2

並列処理をすることで高速化できないでしょうか？ 以下、参考URL

データサイエンスのためのPython入門32〜これだけ知っておけばいいmultiprocessingでの並列処理

Python, Joblibでシンプルな並列処理（joblib.Parallel）

Comments

[bobslay]

ご回答ありがとうございます．
並列化について勉強してみようと思います．

[bobslay]

複数のデータを引数として，同じ関数を何度も呼び出すような処理に対して並列化は有効であるようでしたが，すべてのデータに対してソートを行いたい場合には，うまく並列化する方法は見つかりませんでした．
(データを分割し，それぞれについてソートした後に全体についてソートしていくマージソートのようなことをすると高速化可能？)

Answer 3

実際のデータ量が不明ですが、以下で実用的な速度は出ないでしょうか？
各リストがソート済みであることを生かして、itertools.takewhileにて値の抽出とマージソートを利用しています。

import numpy as np
from heapq import merge
import itertools

a = np.array([1, 2.3, 3.1, 4.6, 6.7])
b = np.array([2.3, 4.2])
c = np.array([])
d = np.array([2.5, 3.2, 4.3, 4.4, 4.8, 5.0])
x = [a, b, c, d]

#x = [np.sort(np.random.random((100,)))*10 for _ in range(100)]

l = [[(v, i) for v in itertools.takewhile( lambda e : e <= 5, ary)] for i, ary in enumerate(x)]
r1, r2 = itertools.tee(merge(*l)) # ソート済みを生かしてマージソートする
result1 = [e[0] for e in r1]
result2 = [e[1] for e in r2]

print(result1)# [1.0, 2.3, 2.3, 2.5, 3.1, 3.2, 4.2, 4.3, 4.4, 4.6, 4.8, 5.0]
print(result2)# [0, 0, 1, 3, 0, 3, 1, 3, 3, 0, 3, 3]

Comments

[bobslay]

ご回答ありがとうございます．
実際に使用したいデータを見たところ，サイズが10^10ほどでした．
そもそもデータ量が大きすぎてこれ以上高速に実行できないのではないかと思ってきました．．．

[can110]

そのサイズだと私の回答の「l = [[(v, i) for v in」での値とインデックスの組を作る部分で、すでに時間がかかりすぎてしまいますね。