グループリストの要素は別のリストに基づいています

Question

私は2つのリスト：inpとbaseを持っています。 inpの各商品をoutのbase変数に基づいてリストに追加したいと思います。 このコードは正常に動作します：

from pprint import pprint as print 

num = 3 
out = [[] for i in range(num)] 
inp = [[1,1],[2,1],[3,2],[7,11],[9,99],[0,-1]] 
base = [0,1,0,2,0,1] 

for i, num in enumerate(base): 
    out[num].append(inp[i]) 
print(out,width=40) 

結果：

[[[1, 1], [3, 2], [9, 99]], 
[[2, 1], [0, -1]], 
[[7, 11]]] 

私はこの使用してnumpyのモジュール（np.arrayとnp.appendかなどを）やりたいです。誰でも助けてくれますか？

Answer 1

numpyの配列としてbaseとinpと仮定すると、我々はこのような何かを行うことができます -

# Get sorted indices for base 
sidx = base.argsort() 

# Get where the sorted version of base changes groups 
split_idx = np.flatnonzero(np.diff(base[sidx])>0)+1 
# OR np.unique(base[sidx],return_index=True)[1][1:] 

# Finally sort inp based on the sorted indices and split based on split_idx 
out = np.split(inp[sidx], split_idx) 

---

が、それはリストのために働くようにするには、我々はそのために我々はできる、いくつかの調整、主にインデックス部分を必要としますnp.takeを使用して、前の方法で列挙した配列にインデックスを置き換えます。したがって、変更されたバージョンは次のようになります：

sidx = np.argsort(base) 
split_idx = np.flatnonzero(np.diff(np.take(base,sidx))>0)+1 
out = np.split(np.take(inp,sidx,axis=0), split_idx)