You can use convolutions. Convolution does something like this (more info here):
It will go through the matrix multipling your filter or pad by the elemnts of the matrix and then adding them together in this case.
For this question, lets first add a new element f to the dataframe so at least one row have more than one element.
>> positions
pos mcap
a 1 1
b 2 4
c 3 3
d 4 2
e 5 5
f 3 2
The positions can be also seen as:
df = pd.crosstab(positions['pos'], positions['mcap'],
values=positions.index, aggfunc=sum)
df
mcap 1 2 3 4 5
pos
1 a NaN NaN NaN NaN
2 NaN NaN NaN b NaN
3 NaN f c NaN NaN
4 NaN d NaN NaN NaN
5 NaN NaN NaN NaN e
df_ones = df.notnull() * 1
mcap 1 2 3 4 5
pos
1 1 0 0 0 0
2 0 0 0 1 0
3 0 1 1 0 0
4 0 1 0 0 0
5 0 0 0 0 1
We can create a window that slides through df_ones and sum the number of elements that fall under the window. This is called 'convolution' (or correlation).
Now let us create a window that avoids the upper left corner element (so it is not counted) and convolute it with our df_ones to get the result:
pad = np.ones_like(df.values)
pad[0, 0] = 0
pad
array([[0, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1]], dtype=object)
counts = ((signal.correlate(df_ones.values, pad,
mode='full')[-df.shape[0]:,
-df.shape[1]:]) * \
df_ones).unstack().replace(0, np.nan).dropna(
).reset_index().rename(columns={0: 'count'})
mcap pos count
0 1 1 5.0
1 2 3 3.0
2 2 4 1.0
3 3 3 1.0
4 4 2 1.0
positions.reset_index().merge(counts,
how='left').fillna(0
).sort_values('pos').set_index('index')
pos mcap count
index
a 1 1 5.0
b 2 4 1.0
c 3 3 1.0
f 3 2 3.0
d 4 2 1.0
e 5 5 0.0
All in a function:
def count_upper(df):
df = pd.crosstab(positions['pos'], positions['mcap'],
values=positions.index, aggfunc=sum)
df_ones = df.notnull() * 1
pad = np.ones_like(df.values)
pad[0, 0] = 0
counts = ((signal.correlate(df_ones.values, pad,
mode='full')[-df.shape[0]:,
-df.shape[1]:]) * df_ones)
counts = counts.unstack().replace(0, np.nan).dropna(
).reset_index().rename(columns={0: 'count'})
result = positions.reset_index().merge(counts,
how='left')
result = result.fillna(0).sort_values('pos').set_index('index')
return result
For your example the result would match your expected result:
positions = pd.DataFrame({"pos" : [1, 2, 3, 4, 5],
"mcap" : [1, 4, 3, 2, 5]},
index = ["a", "b", "c", "d", "e"])
>> count_upper(positions)
pos mcap count
index
a 1 1 4.0
b 2 4 1.0
c 3 3 1.0
d 4 2 1.0
e 5 5 0.0