Hi @phungleson thanks for the contribution!

I think it was a great idea to raise this PR and get a conversation started 👍. I'm a volunteer. I am happy to keep an eye out for comments and changes, review/discuss ideas and generally cheer you on.

It certainly looks like this code would work, but I agree we need to get to a point where we are all happy it's the right implementation.

I suggest the next step in achieving that is to wrap some tests around it so we can experiment with the implementation without risk of breaking things. May not need to be deep coverage at this stage, but it's up to you. I would think about relevant ndarray shapes, types and error cases.

From quick review it looks like its run-time complexity is O(2n) big-oh of 2n. If that's the case, then my instinct is we could improve on that.

Have you had a look for ideas from existing implementations in other technologies? For example numpy and Julia - how do they implement this? Are there any special algorithms people use that are better than the intuitive approach?

Thanks. I hope this helps.