Can you select a variable by binning another variable in a vectorized way?

Problem

I have several variables x that I want to sort in a variable binned_list using some bins.

As an example, x is a random vector in two components, going from 0 to 10/sqrt(2), that I want to sort in the list binned_list by the modulus of x. I have three bins for the modulus: [0, 3.33), [3.33, 6.66) and [6.66, 10) and I want to save different iterations of x into binned_list, which is a list of 3 lists, each one corresponding to values of the modulus of x on that bin.

I can do it in the following way:

import numpy as np
N_bins = 3
Bins = np.linspace(0, 10, N_bins+1)

binned_list = [[] for b in range(N_bins)]

N_elements = 5

np.random.seed(1)

for k in range(3):
    x = np.random.random((N_elements,2))/np.sqrt(2)*10
    mod_x = np.sqrt(x[:,0]**2 + x[:,1]**2)
    dig_x = np.digitize( mod_x, bins = Bins ) - 1

    for y in range(len(x)):
        binned_list[dig_x[y]].append(x[y])

Output:

[[array([8.08752089e-04, 2.13781412e+00]),
  array([1.03772086, 0.65293247]),
  array([1.31705859, 2.44348333]),
  array([0.99268556, 1.40078906]),
  array([0.60135339, 0.27615902])],
 [array([2.94879087, 5.09346334]),
  array([2.80556972, 3.81000966]),
  array([2.96415284, 4.84523355]),
  array([1.44569572, 6.20922794]),
  array([0.19365953, 4.74092123]),
  array([2.95079056, 3.95053366]),
  array([2.21624362, 4.89546016]),
  array([1.20088241, 6.20940519])],
 [array([5.66211915, 6.84664326]), array([6.19700713, 6.32582438])]]

Question

Once I have digitized the elements of x, can I avoid looping through them in order to save them in variable binned_list? I would like to do this in a vectorized way in order to make the code more efficient.

I thought of something like:

binned_list[dig_x].append(x)

But I can't slice a list with an array. Also if I define binned_list as an array I can't append either.

You can avoid the nested for loop and only loop through the bins by using masking with nonzero()

x = np.random.random((N_elements,2))/np.sqrt(2)*10
mod_x = np.sqrt(np.sum(x**2,1))
dig_x = np.digitize(mod_x,bins=Bins)-1
for i in range(N_bins):
    binned_list[i] = x[(dig_x==i).nonzero()]

I compared @mpw2 answer and mine and his answer is indeed a bit faster, around a factor 1.5-5 faster for different number of iterations and number of elements that I tried:

import numpy as np
import time

N_bins = 3
Bins = np.linspace(0, 10, N_bins+1)

binned_list1 = [[] for b in range(N_bins)]
binned_list2 = [[] for b in range(N_bins)]


N_elements = 10000

np.random.seed(1)

k_iter = 10

x = np.random.random((k_iter, N_elements,2))/np.sqrt(2)*10


start_time = time.time()
for k in range(k_iter):
    mod_x = np.sqrt(x[k,:,0]**2 + x[k,:,1]**2)
    dig_x = np.digitize( mod_x, bins = Bins ) - 1

    for y in range(len(x[k])):
        binned_list1[dig_x[y]].append(x[k,y])

print("1: %s s" % (time.time() - start_time))

start_time = time.time()
for k in range(k_iter):
    mod_x = np.sqrt(x[k,:,0]**2 + x[k,:,1]**2)
    dig_x = np.digitize( mod_x, bins = Bins ) - 1

    for i in range(N_bins):
        binned_list2[i].extend(x[k,(dig_x==i).nonzero()[0]])

print("2: %s s" % (time.time() - start_time))

for i in range(N_bins):
    print(np.allclose(np.array(binned_list1[i]), np.array(binned_list2[i])))

Output

1: 0.039897918701171875 s
2: 0.01396489143371582 s
True
True
True