group list elements based on another list
<p>I have two list: <code>inp</code> and <code>base</code>. i want to add each item in <code>inp</code> to a list in <code>out</code> based on position in <code>base</code> variable.
this code works fine:</p>

<pre><code>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)
</code></pre>

<p>result:</p>

<pre><code>[[[1, 1], [3, 2], [9, 99]],
 [[2, 1], [0, -1]],
 [[7, 11]]]
</code></pre>

<p>i want to do this using <strong><em>numpy</em></strong> module(np.array and np.append or etc.). Can anyone help me?</p>

<p>Assuming <code>base</code>and <code>inp</code> as NumPy arrays, we could do something like this -</p>

<pre><code># 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])&gt;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)
</code></pre>

<hr>

<p>To make it work for lists, we need few tweaks, mainly the indexing part, for which we can use <code>np.take</code> to replace the indexing into arrays as listed in the earlier approach. So, the modified version would be -</p>

<pre><code>sidx = np.argsort(base)
split_idx = np.flatnonzero(np.diff(np.take(base,sidx))&gt;0)+1
out = np.split(np.take(inp,sidx,axis=0), split_idx)
</code></pre>

python arrays numpy