Consider the following pd.DataFrame

import numpy as np
import pandas as pd

start_end = pd.DataFrame([[(0, 3), (4, 5), (6, 12)], [(7, 10), (11, 90), (91, 99)]])
values = np.random.rand(1, 99)

The start_end is a pd.DataFrame of shape (X, Y) where each value inside is a tuple of (start_location, end_location) in the values vector. Another way of saying that the values in a particular cell is a vector of different lengths.


If I want to find the mean (for example) of the vector values for each of the cells in the pd.DataFrame, how can I do this in a cost effective way?

I managed to achieve this with an .apply function, but it's quite slow.

I guess I need to find some way to present it in a numpy array and then map it back to the 2d data-frame, but I can't figure out how.


  • The distance between start end can varies and outliers can exist.
  • The cell start / end is always non-overlapping with the other cells (it will be interest to see if this prerequisite affect speed of solution).

The generalized problem

More generally speaking I this as a recurring problem of how to make a 3d array, where one of the dimensions is not of equal length to a 2d matrix via some transformation function (mean, min, etc.)

  • What are the sizes of your data? – Quang Hoang Jul 7 at 13:16
  • Usually it's about 10,000 x 100,000 (the start_end) and the vector is about 10,000,000 If you think - the vector is too small, this is because often the start_end has tuples of negative or nan values (i.e. this cell does not contain info in the values vector) – Newskooler Jul 7 at 13:20
  • Shouldn't start_end be a 2D array of (N,2) shape? – Divakar Jul 7 at 13:21
  • For each row in start_end, are the cells non-overlapping and in increasing order? – Quang Hoang Jul 7 at 13:22
  • @Divakar I updated the description. Maybe an (N, 2) would be the state after one transformation step, as long as it can be mapped back to the 2d matrix of original size. – Newskooler Jul 7 at 13:26

Prospective approach

Looking at your sample data :

In [64]: start_end
         0         1         2
0   (1, 6)    (4, 5)   (6, 12)
1  (7, 10)  (11, 12)  (13, 19)

It is indeed non-overlapping for each row, but not across the entire dataset.

Now, we have np.ufunc.reduceat that gives us ufunc reduction for each slice :

ufunc(ar[indices[i]: indices[i + 1]])

as long as indices[i] < indices[i+1].

So, with ufunc(ar, indices), we would get :

[ufunc(ar[indices[0]: indices[1]]), ufunc(ar[indices[1]: indices[2]]), ..]

In our case, for each tuple (x,y), we know x<y. With stacked version, we have :


If we flatten, it would be :

[x1,y1,x2,y2,x3,y3, ...]

So, we might not have y1<x2, but that's okay, because we don't need ufunc reduction for that one and similarly for the pair : y2,x3. But that's okay as they could be skipped with a stepsize slicing of the final output.

Thus, we would have :

# Inputs : a (1D array), start_end (2D array of shape (N,2))
lens = start_end[:,1]-start_end[:,0]
out = np.add.reduceat(a, start_end.ravel())[::2]/lens

np.add.reduceat() part gives us the sliced summations. We needed the division by lens for the average computations.

Sample run -

In [47]: a
array([0.49264042, 0.00506412, 0.61419663, 0.77596769, 0.50721381,
       0.76943416, 0.83570173, 0.2085408 , 0.38992344, 0.64348176,
       0.3168665 , 0.78276451, 0.03779647, 0.33456905, 0.93971763,
       0.49663649, 0.4060438 , 0.8711461 , 0.27630025, 0.17129342])

In [48]: start_end
array([[ 1,  3],
       [ 4,  5],
       [ 6, 12],
       [ 7, 10],
       [11, 12],
       [13, 19]])

In [49]: [np.mean(a[i:j]) for (i,j) in start_end]

In [50]: lens = start_end[:,1]-start_end[:,0]
    ...: out = np.add.reduceat(a, start_end.ravel())[::2]/lens

In [51]: out
array([0.30963037, 0.50721381, 0.52954646, 0.413982  , 0.78276451,

For completeness, referring back to given sample, the conversion steps were :

# Given start_end as df and values as a 2D array
start_end = np.vstack(np.concatenate(start_end.values)) 
a = values.ravel()  

For other ufuncs that have reduceat method, we will just replace np.add.reduceat

| improve this answer | |
  • can you please show the conversion step form the pd.DataFrame I have to the (N, 2) array? I struggle with this part. – Newskooler Jul 7 at 13:39
  • @Newskooler Added in the post. – Divakar Jul 7 at 13:42
  • 2
    @Newskooler Well you have values = np.random.rand(1, 99). So its not a vector, but a 2D array. We need a 1D array for processing, hence flattening with .ravel(). If that's ineeded a vector (1D), you can skip ravel. – Divakar Jul 7 at 13:46
  • 1
    @Newskooler Note that I changed the first pair to [ 1, 3] to bring variety there. – Divakar Jul 7 at 13:56
  • 1
    @Newskooler Only way to not skip is when the end of one slice is same as start of the next slice. Then, we can just use np.add.reduceat(a, start_end[:,0]). – Divakar Jul 7 at 14:40

For computing mean in your case, you'll never go as fast as if you precompute cumulative sums first using numpy.cumsum for instance. Check out the following code:

import numpy as np
import pandas as pd
import time

R = 1_000
C = 10_000
M = 100

# Generation of test case
start = np.random.randint(0, M-1, (R*C,1))
start = np.where(np.logical_and(start>=end, end>1), end-1, start)
end = np.where(np.logical_and(start>=end, start<M-1), start+1, end)
start_end = np.hstack((start, end))

values = np.random.rand(M)

t_start = time.time()
# Basic mean dataframe
mean = np.add.reduceat(values, start_end.ravel())[::2]/lens
print('Timre 1:', time.time()-t_start, 's')

t_start = time.time()
#Cumulative sum
cum_values = np.zeros((values.size+1,))
cum_values[1:] = np.cumsum(values)
# Compute mean dataframe
mean_2 = (cum_values[start_end[:,1]]-cum_values[start_end[:,0]])/(start_end[:,1]-start_end[:,0])
print('Timre 2:', time.time()-t_start, 's')

print('Results are equal!' if np.allclose(mean, mean_2) else 'Results differ!')
print('Norm of the difference:', np.linalg.norm(mean - mean_2))


% python3 script.py
Timre 1: 0.48940515518188477 s
Timre 2: 0.16983389854431152 s
Results are equal!
Norm of the difference: 2.545241707481022e-12

The difference in performance gets even worse when M increases. For M=5000 you get:

% python3 script.py
Timre 1: 4.5356669425964355 s
Timre 2: 0.1772768497467041 s
Results are equal!
Norm of the difference: 1.0660592585125616e-10
| improve this answer | |
  • Looks promising indeed, good idea with cumsum. – Divakar Jul 7 at 16:07
  • Yes! After computing cumsum with complexity O(N), all the other mean can be computed with O(1). It's very powerful especially when you have to compute rolling window analysis. You can even use this trick to compute standard deviations or correlations in O(1). – bousof Jul 7 at 16:20
  • Can be numerically unstable, though. – Paul Panzer Jul 7 at 17:15
  • You mean if you get an overflow from cumsum ? – bousof Jul 7 at 17:19
  • Even before that if your increments are small compared to the running total you can lose a significant number of bits: Example, just to show the principle a=np.array([2.0**40,1.2345678]) np.diff(a.cumsum()) yields array([1.23461914])`. – Paul Panzer Jul 8 at 3:19

Your Answer

