python 二次元データ ソート

Question

pythonで次のコードのようになっている二次元のジャグ配列のデータを小さい順に取り出したいです．
また取り出した要素の添字も取り出したデータと同じ順番で並べたものも取得したいです．
次のコードでは，二次元配列はndarrayのリストになっていますが，単純な多重リストの方が速く実装できるのであれば，そちらの方法でも構いません．

import numpy as np

a = np.array([0.4, 0.8, 2.3])
b = np.array([0.3, 0.6])
c = np.array([2, 2.1, 3.4, 4.4])
d = np.array([0.4, 0.7])

abcd = [a, b, c, d]

〜〜〜ここでソートを実行〜〜〜

得たい結果
result1 = [0.3, 0.4, 0.6, 0.7, 0.8, 2, 2.1, 2.3, 3.4, 4.4]
result2 = [[1], [0, 3], [1], [3], [0], [2], [2], [0], [2], [2]]

実際に使用する際にはさらに大きなデータで実行したいと考えているので，できるだけ高速なコードが知りたいです．お願いいたします．

Comments

[退会済みユーザー]

「得たい結果」の２つ目のリスト result2 の二番目の要素 [0, 3] は何を表しているのでしょうか？

[bobslay]

0.4という値が0番目(a)・3番目(b)どちらにもあるので，同率ということで[0,3]という書き方をしています．わかりにくく申し訳ありません．

[退会済みユーザー]

ご回答ありがとうございます。なるほど、そういう意味ですね。了解しました。

Answer 1

numpyでやってみました。
argsortしたあとuniqueという似たような処理を2回やっているので、ちょっと無駄が多いかもしれないのですが、同率の組み合わせを作るためにやっています。
pandasのgroupbyや、itertools.groupbyでもできると思います。

import numpy as np

a = np.array([0.4, 0.8, 2.3])
b = np.array([0.3, 0.6])
c = np.array([2, 2.1, 3.4, 4.4])
d = np.array([0.4, 0.7])

abcd = [a, b, c, d]

arr = np.concatenate(abcd)
group = np.concatenate([[i] * len(x) for i, x in enumerate(abcd)])

sort_idx = np.argsort(arr)
arr_sorted = arr[sort_idx]
grp_sorted = group[sort_idx]

result1, uniq_idx = np.unique(arr_sorted, return_index=True)
result2 = np.split(grp_sorted, uniq_idx[1:])

print(result1)
# [0.3 0.4 0.6 0.7 0.8 2.  2.1 2.3 3.4 4.4]
print(result2)
# [array([1]), array([0, 3]), array([1]), array([3]), array([0]), array([2]), array([2]), array([0]), array([2]), array([2])]

Comments

[kirara0048]

group = np.arange(len(abcd)).repeat([*map(len, abcd)])
のほうが良いかもしれません。

[bsdfan]

kirara0048さん、ありがとうございます。
arangeとrepeatを組み合わせて書けますね。思いつきませんでした。勉強になります。

Answer 2

これでどうでしょう？

import numpy as np
from collections import defaultdict


a = np.array([0.4, 0.8, 2.3])
b = np.array([0.3, 0.6])
c = np.array([2, 2.1, 3.4, 4.4])
d = np.array([0.4, 0.7])
abcd = [a, b, c, d]

data_map = defaultdict(list)

for idx, ls in enumerate(abcd):
  for v in ls:
    data_map[v].append(idx)


result1 = sorted(data_map.keys())
result2 = [data_map[v] for v in result1]

print(result1)  # => [0.3, 0.4, 0.6, 0.7, 0.8, 2.0, 2.1, 2.3, 3.4, 4.4]
print(result2)  # => [[1], [0, 3], [1], [3], [0], [2], [2], [0], [2], [2]]

Answer 3

以下です。

>>> a = [0.4, 0.8, 2.3]
>>> b = [0.3, 0.6]
>>> c = [2, 2.1, 3.4, 4.4]
>>> d = [0.4, 0.7]
>>>
>>> abcd = [a, b, c, d]
>>>
>>> all_set = {}
>>> for i, lst in enumerate(abcd):
...     for x in lst:
...         if x in all_set:
...             all_set[x].append(i)
...         else:
...             all_set[x] = [i]
...
>>> all_list = sorted(all_set.items(), key=lambda item: item[0])
>>>
>>> result1 = [x for x, y in all_list]
>>> result2 = [y for x, y in all_list]
>>> print(result1)
[0.3, 0.4, 0.6, 0.7, 0.8, 2, 2.1, 2.3, 3.4, 4.4]
>>> print(result2)
[[1], [0, 3], [1], [3], [0], [2], [2], [0], [2], [2]]

追加回答

なるべく、機械命令で実行する部分を増やすために、numpyとpandasを使った方法です。
numpyを使っているため、「0.4という値が0番目(a)・3番目(b)どちらにもあるので，同率ということで[0,3]という書き方」はやっていません。それをやると遅くなります。

>>> import numpy as np
>>> import pandas as pd
>>>
>>> a = np.array([0.4, 0.8, 2.3])
>>> b = np.array([0.3, 0.6])
>>> c = np.array([2, 2.1, 3.4, 4.4])
>>> d = np.array([0.4, 0.7])
>>>
>>> abcd = [a, b, c, d]
>>>
>>> num_elements = sum([len(e) for e in abcd])
>>> all_array = np.zeros(2 * num_elements).reshape(num_elements, 2)
>>>
>>> start = 0
>>> for i, a in enumerate(abcd):
...     end = start + len(a)
...     all_array[start:end, 0] = a
...     all_array[start:end, 1] = i
...     start = end
...
>>> df = pd.DataFrame(all_array)
>>> df.sort_values(0, inplace=True)
>>> print(df)
      0    1
3   0.3  1.0
0   0.4  0.0
9   0.4  3.0
4   0.6  1.0
10  0.7  3.0
1   0.8  0.0
5   2.0  2.0
6   2.1  2.0
2   2.3  0.0
7   3.4  2.0
8   4.4  2.0
>>> result1 = df[0].values
>>> result2 = df[1].values.astype(int)
>>> print(result1)
[0.3 0.4 0.4 0.6 0.7 0.8 2.  2.1 2.3 3.4 4.4]
>>> print(result2)
[1 0 3 1 3 0 2 2 0 2 2]

最初の回答よりは速いと思います。

Comments

[bobslay]

ご回答ありがとうございます．
ただいま実際に使用したいデータに対して実行してみたところ，やはりndarrayをリストに変換するのに時間がかかってしまいました．
tolist()関数を用いて二次元配列の内側の要素を一つずつ変換しているので遅いと思うのですが，何か他に良い方法はあるでしょうか？
もしくは，やはり質問文のコードのデータ形式のまま高速にソートする方法はありますでしょうか？

[ppaul]

ndarrayをリストに変換しなくても、同じコードで動きます。
しかし、それほど速くはならないでしょう。
ちょっと考えてみますが、numpyだけでできるかというと難しそうです。