That’s a nice improvement
But there’s still a classic mistake:
Here, you create an array on the right hand side, and then copy the elements over. It’s better to never create that rhs array. Just loop over i and j and directly write the values into res.