Rohin pointed out that the input matrix 0, 1, …, 15 is not a safe input for Part 2 of the homework. This is an excellent catch, Rohin.
If you already submitted your homework, don’t worry about this error. If you have not, change your input to 1, 2, …, 16.
As an extra credit, explain why the 0..15 input suffices for Part 1 but fails for Part 2.
--Ras