MCO16406 - Virus Attack - Getting WA with SQRT decomposition

Hello, I’ve been attempting this question - using sqrt decomposition but keep getting WA in task 1 itself. I’m not able to figure out where I’m going wrong or which test case I’m missing. Can anyone please help me with this?

My code -