k-й наименьший элемент таблицы умножения

Задача. Дана таблица умножения размером m * n и целое число k, нужно найти the k-й элемент таблицы умножения, если взять все значения таблицы и отсортировать их в возрастающем порядке.

Пример 1:

Дано: m = 3, n = 3, k = 5
Результат: 3
Пояснение: 
Таблица умножения 3х3:
1 2 3
2 4 6
3 6 9

5-й наименьший элемент - это 3 (1, 2, 2, 3, 3, 4, 6, 6, 9)

Пример 2:

Дано: m = 2, n = 3, k = 6
Результат: 6
Пояснение: 
Таблица умножения 2х3:
1 2 3
2 4 6

6-й наименьший элемент - это 6 (1, 2, 2, 3, 4, 6).

Ограничения:

1. m иn из интервала [1, 30000].
2. k из интервала [1, m * n]

Брутфорс-идея: берем все элементы матрицы (предварительно их посчитав), сортируем и берем k-й элемент. Анализ сложности:

• Сложность по времени: $O(m*n)$ на подсёт элементов таблицы, и $O(m*n\log(m*n))$$\ast \log (m\ast n))$ на их сортировку.
• Затраты на память: $O(m*n)$.

Оптимизация: мы знаем наименьший элемент (1) и наибольший элемент (m*n) таблицы. Если мы возьмём какое-либо Х число из интервала между ними, то для этого числа нам нужно проверить:

• какие числа из таблицы умножения меньше Х;
• сколько таких чисел.

Допустим их z. Если z>k, то есть Х находится на позиции z (+= количество его повторений) в отсортированном массиве, то нам нужно уменьшить число Х и проверить новое значение на соотношение z и k. Обратная ситуация: если z<k, то нам нужно увеличить число Х и опять проверить новое значение на соотношение z и k. В итоге нам нужно выбрать такое число R(езультат), чтобы оно находилось на k-й позиции.

Эти уменьшения/увеличения числа Х наводят на мысль  об использовании двоичного поиска (дихотомия).

Компилируемое решение

int findKthNumber(int m, int n, int k) {
    // верхняя и нижняя границы входных данных
    int lo = 1; int hi = m*n;
    while(lo < hi) {
        // берем число посередине
        int mid = lo + (hi - lo) / 2;
        // нам нужно посчитать, скольо элементов в матрице
        // имеют значение меньше либо равное mid
        int count = 0,  j = m;
        // просматриваем все столбцы от 1 до n
        for(int i = 1; i <= n; i++) {
            // очевидно, чем раньше столбец, 
            // тем позже (ниже) у него будут числа, 
            // которые превосходят нашу опорную точку mid;
            // поэтому просматриваем столбец снизу вверх,
            // чтобы найти последний первый индекс, для которого
            // не выполнено условие i*j <= mid;
            // эти элементы мы пропускаем, "укорачивая" индекс j
            // И для следующего столбца стартуем с "укороченного" индекса j,
            // поскольку для столбца с бОльшим индексом i, очевидно, 
            // i*j будут больше mid.
            // Например, если для столбца 1 мы укоротили j с m до p, то есть
            // (p+1)*1>mid, то очевидно, что (p+1)*2>mid тоже.
            while(j >=1 && i*j > mid) j--; // укоротили j
            // здесь j наибольший в столбце i, что i*j<=mid
            count += (j); // добавили в счётчик j элементов столбца i
        }
        // проверяем количесвто элементов таких, что они меньше mid
        // если это количество меньше заданного k, то mid (и всё, что меньше его)
        // слишком маленькое,
        // поэтому поднимаем нижнюю границу
        if(count < k) lo = mid + 1;
        // если это количество больше либо равно заданного k, то
        // верхнюю границу нужно опустить, но включить mid, 
        // поскольку для него может выполняться условие count==k
        else hi = mid;
    }
    // lo == hi, больше двигаться некуда, возвращаем lo или hi
    return hi; // return lo; тоже работает
}

 Почему без разницы, что писать в return?

Допустим, lo!=hi. Поскольку мы вышли из цикла while, то lo > hi. Два случая, когда меняюися lo и hi

• newhi = mid, после чего lo>newhi, значит на предыдущем шаге было lo>mid (потому что lo не поменялось, но получилось, что lo>hi, которое=mid). А hi было ещё больше (иначе мы бы вышли из цикла ранее). То есть lo = mid + a, hi = mid + b, 0<a<b. Но тогда подсчёт среднего значения lo + (hi-lo)/2 = (lo + hi)/2 = (mid + a + mid + b)/2 = mid + (a+b)/2. А это больше mid. То есть мы получаем противоречие. Ситауция невозможна.
• newlo = mid + 1, после чего newlo>hi, значит на предыдущем шаге были такие lo и hi, что (lo+hi)/2+1>hi (hi не поменялось). Или (hi-lo)<2, то есть это два элемента рядом: e и e+1. Для них тогда mid = e (округление по нижней границе). Но тогда newlo=mid+1=e+1=hi. То есть (new)lo=hi, что и требовалось доказать.

Почему return возваращает элемент из матрицы умножения? А что если lo=hi - это элемент, которого нет в матрице? 

Допустим, в ответе return у нас некое число Х, которое не присутствует в матрице. Для этого числа мы посчитали count, то есть количесво элементов в исходной матрице, которые меньше, чем X. Мы возьмём максимальный среди них элемент M, и для него count будет точно таким же, как и для X. То есть count - это количество элементов, кторые меньше, либо равны М. А поскольку М принадлежит матрице, и при этом M<X, то получается, что X - это не наименьший элемент, как того требует задача, поэтому мы вернем M.

Интересное примечание.
Можно оптимизировать строку

while(j >=1 && i*j > mid) j--; // укоротили j
count += (j); // добавили в счётчик j элементов столбца i

и представить, как: count += Math.max(mid/i, m). Ведь mid/i даёт нам наибольший индекс j, для которого i*j<=mid. Но почему-то такой код рабоает МЕДЛЕННЕЕ, чем перебор по j, который у нас.

Runtime: 8 ms (в "оптимизированном" виде 11 ms), faster than 89.84% of Java online submissions for Kth Smallest Number in Multiplication Table.
Memory Usage: 33.1 MB, less than 33.33% of Java online submissions for Kth Smallest Number in Multiplication Table.

Анализ сложности

• Сложность по времени: $O(m * \log (m*n))$. Наш двоичный поиск делит интервал $\text{[lo, hi]}$ на две равные части на каждом шаге. На каждом шаге мы делаем цикл for за время $O(m)$, если рассмотреть "оптимизированный" вариант.
• Память: $O(1)$. Мы держим в памяти только lo, hi, mid, count, i, j, все целочисленные, то есть 6 ячеек памяти для целочисленных значений зарезервировано.