python version of Xian’s code

John Haman

2020/09/09

Categories: R Python Tags: R Python

Sort of inspired to find a faster solution to a Riddler by Xian, I rewrote some of his code in Python a got a ~10-15x speed increase. Read Xian’s post for a link to the original Riddle.

Load Libs

library(reticulate)
library(tictoc)

R Solution

tic()
simz=t(apply(matrix(runif(3*1e5),ncol=3),1,sort))
mean((simz[,1]>.5)*simz[,1]+
       (simz[,1]<.5)*(simz[,2]>.5)*(simz[,2]-simz[,1])+
       (simz[,2]<.5)*(simz[,3]>.5)*(simz[,3]-simz[,2])+
       (simz[,3]<.5)*(1-simz[,3]))
## [1] 0.4685494
toc()
## 4.123 sec elapsed

Python solution

import numpy as np
import time

start = time.time()

mat = np.random.uniform(size=3 * int(1e5)).reshape((int(1e5), 3))
simz = np.apply_along_axis(np.sort, 1, mat)

a = np.mean(
    (simz[:, 0] > 0.5) * simz[:, 0]
    + (simz[:, 0] < 0.5) * (simz[:, 1] > 0.5) * (simz[:, 1] - simz[:, 0])
    + (simz[:, 1] < 0.5) * (simz[:, 2] > 0.5) * (simz[:, 2] - simz[:, 1])
    + (simz[:, 2] < 0.5) * (1 - simz[:, 2]))

end = time.time()

print(end - start)
## 0.5366349220275879

I’m a Python noob, but this seems pretty great :)