title: Partitioning the rows of a numpy array into buckets
body: <p>I am writing code to process a large point cloud. I have been able to successfully vectorize most of my code to make it efficient, however, I cannot think of a good way of achieving this:</p>
  
  <p>I have an nx3 numpy array that represents points in 3d (x,y,z) with z >= 0. I want to create a list of length k, which is produced by bucketing the rows of the numpy array by z value. i.e. if a point has z value $z_i$, then this point should go into bucket $j =\lfloor z_i / d \rfloor$, where $d$ is the resolution of the buckets. $d$ is precomputed as $z_{max} / k$, where k is the fixed number of buckets. Right now, I have tried the following:</p>
  
  <pre><code>buckets = [set() for _ in range(num_buckets)]
  d = max_z / num_slices
  
  for i in range(len(points)):
      p = points[i]
      k = int(math.floor(p[2] / d))
      buckets[min(k, num_buckets - 1)] |= {tuple(p)}
  </code></pre>
  
  <p>However, this is very inefficient since I am iterating over the whole dataset of points (which is very large). On the other hand, it emphasizes the independence on the number of buckets. I have closed form for the index of a point given its z value, so I should not need to iterate over the bucket list.</p>
  
  <p>That being said, num_buckets is much smaller than the number of points, so I also wrote this, which is must faster:</p>
  
  <pre><code>for i in range(num_buckets):
      ps = points[np.floor(in_fov[:,2] / d) == i]
      buckets[i] = ps
  </code></pre>
  
  <p>This is certainly faster, but now the runtime depends on num_buckets, and the growth is fairly troubling since I am doing a numpy filtering on a large array for every bucket.</p>
  
  <p>My question is if there is a better way to do this that allows me to allocate each row of this nx3 array to a list index that allows me to leverage the fact that I can compute the index for each point independently, and hopefully eliminates the need to iterate over num_buckets.</p>