I came up with the same logic as yours, but whenever if found arr[j]==2 or arr[j]==0, I directly pushed them into the return array. This is why my code failed on tests cases like [2,2]. It is important to note that you just store the values in temp variables (a and b) and then construct the return array later.
@claytonjwong Yeah I've known n(n+1)/2 for about 30 years and used it many times. Here I chose not to. Five of my twelve solutions that I have in the editor actually use it, but I like the above solution best.
Your test is indeed bad. Times this small aren't trustworthy. Especially since you're even printing during the timing. Printing is slow. You might effectively measure the printing instead of the calculation. Better make it take at least one second or so, by using a larger list or repeating the calculation often enough. And do just calculation, no printing (during the timing).
Btw, you can simply do nums = list(range(1, n+1)) instead of that loop.
Or if you want to confuse/educate people a bit, do nums, nums = list(range(1, n+1)), 3 :-)