Here's a really crazy way to do it, which involves sorting and indexing rather than adding a new dimension. This is sort of like the sort-based method used by np.unique.
First find the sorted indices in each row:
rows = np.repeat(np.arange(x.shape[0]), x.shape[1]) # [0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3]
cols = np.argsort(x, axis=1).ravel() # [0, 2, 1, 2, 1, 0, 0, 1, 2, 1, 0, 2]
Now you can create an array of sorted elements per-column, both unweighted and weighted. The former will be used to get the indices for summing, the latter will actually be summed.
u = x[rows, cols] # [1, 1, 2, 1, 2, 3, 1, 2, 2, 1, 3, 3]
v = np.broadcast_to(w, x.shape)[rows, cols] # [0.3, 0.3, 0.4, 0.3, 0.4, 0.3, 0.3, 0.4, 0.3, 0.4, 0.3, 0.3]
You can find the breakpoints to apply np.add.reduce at:
row_breaks = np.diff(rows).astype(bool) # [0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0]
col_breaks = np.diff(u).astype(bool) # [0, 1, 1, 1, 1, 1, 1, 0, 1, 1, 0]
break_mask = row_breaks | col_breaks # [0, 1, 1, 1, 1, 1, 1, 0, 1, 1, 0]
breaks = np.r_[0, np.flatnonzero(break_mask) + 1] # [ 0, 2, 3, 4, 5, 6, 7, 9, 10]
Now you have the sums of the weights for identical numbers in each row:
sums = np.add.reduceat(v, breaks) # [0.6, 0.4, 0.3, 0.4, 0.3, 0.3, 0.7, 0.4, 0.6]
But you need to break them up into segments corresponding to the number of unique elements per row:
unique_counts = np.add.reduceat(break_mask, np.arange(0, x.size, x.shape[1]))
unique_counts[-1] += 1 # The last segment will be missing from the mask: # [2, 3, 2, 2]
unique_rows = np.repeat(np.arange(x.shape[0]), unique_counts) # [0, 0, 1, 1, 1, 2, 2, 3, 3]
You can now sort each segment to find the maximum value:
indices = np.lexsort(np.stack((sums, unique_rows), axis=0)) # [1, 0, 2, 4, 3, 5, 6, 7, 8]
The index at the end of each run is given by:
max_inds = np.cumsum(unique_counts) - 1 # [1, 4, 6, 8]
So the maximum sums are:
sums[indices[max_inds]] # [0.6, 0.4, 0.7, 0.6]
And you can unravel the indices-within indices to get the correct element from each row. Notice that max_inds, and everything that depends on it is the same size as x.shape[1], as expected:
result = u[breaks[indices[max_ind]]]
This method does not look very pretty, but it is likely more space efficient than using an extra dimension on the array. Additionally, it works regardless of the numbers in x. Notice that I never subtracted anything or adjusted x in any way. In fact, all the rows are treated independently, and the coincidence of a maximum element being identical to the minimum of the next is broken by row_breaks when constructing breaks.
TL;DR
Enjoy:
def weighted_vote(x, w):
rows = np.repeat(np.arange(x.shape[0]), x.shape[1])
cols = np.argsort(x, axis=1).ravel()
u = x[rows, cols]
v = np.broadcast_to(w, x.shape)[rows, cols]
row_breaks = np.diff(rows).astype(bool)
col_breaks = np.diff(u).astype(bool)
break_mask = row_breaks | col_breaks
breaks = np.r_[0, np.flatnonzero(break_mask) + 1]
sums = np.add.reduceat(v, breaks)
unique_counts = np.add.reduceat(break_mask, np.arange(0, x.size, x.shape[1]))
unique_counts[-1] += 1
unique_rows = np.repeat(np.arange(x.shape[0]), unique_counts)
indices = np.lexsort(np.stack((sums, unique_rows), axis=0))
max_inds = np.cumsum(unique_counts) - 1
return u[breaks[indices[max_inds]]]
Benchmarks
Benchmarks are run in following format:
rows = ...
cols = ...
x = np.random.randint(cols, size=(rows, cols)) + 1
w = np.random.rand(cols)
%timeit weighted_vote_MP(x, w)
%timeit weighted_vote_JG(x, w)
assert (weighted_vote_MP(x, w) == weighted_vote_JG(x, w)).all()
I used the following generalization for weighted_vote_JG, with appropriate corrections:
def weighted_vote_JG(x, w):
i = np.arange(w.size) + 1
m = (x[None, ...] == i.reshape(-1, 1, 1))
return np.argmax(np.sum(m * w, axis=2), axis=0) + 1
Rows: 100, Cols: 10
MP: 440 µs ± 5.12 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
* JG: 153 µs ± 796 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
Rows: 1000, Cols: 10
MP: 2.53 ms ± 43.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
* JG: 1.03 ms ± 12 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
Rows: 10000, Cols: 10
MP: 23.5 ms ± 200 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
* JG: 16.6 ms ± 67.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
Rows: 100000, Cols: 10
MP: 322 ms ± 3.11 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
* JG: 188 ms ± 858 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
Rows: 100, Cols: 100
* MP: 3.31 ms ± 257 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
JG: 12.6 ms ± 244 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
Rows: 1000, Cols: 100
* MP: 31 ms ± 159 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
JG: 134 ms ± 581 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
Rows: 10000, Cols: 100
* MP: 417 ms ± 7.06 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
JG: 1.42 s ± 126 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
Rows: 100000, Cols: 100
* MP: 4.94 s ± 25.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
JG: MemoryError: Unable to allocate 7.45 GiB for an array with shape (100, 100000, 100) and data type float64
Moral of the story: for small number of columns and weights, the expanded solution is faster. For a larger number of columns, use my version instead.