I implemented the solution described in the editorial, with KMP instead of Z-function, but all my submissions ended up TLE41 (261016696) and I couldn't get the log^2 solution right.

By dividing the submission into $$$>\sqrt{n}$$$ and $$$\le\sqrt{n}$$$ parts I found my solution would finish in 3.2 seconds and I tried my best to squeeze in but still couldn't.

Could anyone help me optimize the constant factor further, or is KMP unable to pass?