Wydaje mi się, że za bardzo skomplikowałeś. Może i tak też można, ale mi
wystarczyło tylko samo sortowanie, a potem wyszukiwanie binarne w posortowanym zbiorze.
Podejrzewam, że wiem skąd wziąłeś równanie #4 < #3, ale w zadaniu występują dwie zmienne n i k, a nie jedna i dodatkowo podejrzewam, że czas testu #3 ledwie zmieścił się w limicie?
Więc jedyne co mogę Ci doradzić, to zrób zadanie-kod najprościej jak się da.