Разбор задачи «Lantern Code»

Идея задачи

Нам дано число n. Нужно рассматривать все числа, которые можно представить в виде суммы различных степеней числа n:

n^0, n^1, n^2, n^3, ...

Каждую степень можно взять не более одного раза.

Например, при n = 3 можно получить такие числа:

• 1 = 3^0
• 3 = 3^1
• 4 = 3^0 + 3^1
• 9 = 3^2
• 10 = 3^0 + 3^2
• 12 = 3^1 + 3^2
• 13 = 3^0 + 3^1 + 3^2

Эти числа упорядочиваются по возрастанию. Для каждого запроса нужно найти k-е такое число по модулю 10^9 + 7.

Главное наблюдение

Посмотрим на двоичную запись числа k.

Пусть, например:

k = 13

В двоичной системе:

13 = 1101_2

Это значит:

13 = 2^0 + 2^2 + 2^3

Теперь заметим важную идею: если в двоичной записи k какой-то бит равен 1, то в ответе должна присутствовать соответствующая степень n.

То есть вместо степеней двойки:

2^0, 2^2, 2^3

мы берём степени числа n:

n^0, n^2, n^3

Тогда ответ будет:

n^0 + n^2 + n^3

Именно так строится k-е допустимое число.

Почему это работает

Все допустимые числа имеют вид:

b0 * n^0 + b1 * n^1 + b2 * n^2 + ...

где каждый коэффициент bi равен либо 0, либо 1.

Это полностью похоже на запись числа в двоичной системе:

b0 * 2^0 + b1 * 2^1 + b2 * 2^2 + ...

Получается очень удобное соответствие:

• бит 0 в числе k — степень n^i не берём;
• бит 1 в числе k — степень n^i берём.

Поэтому нужно просто пройти по битам числа k.

Разберём пример

Пусть:

n = 3, k = 5

Запишем 5 в двоичной системе:

5 = 101_2

Значит, берём:

• 3^0
• 3^2

Не берём:

• 3^1

Получаем:

3^0 + 3^2 = 1 + 9 = 10

Значит, ответ равен 10.

И действительно, последовательность допустимых чисел для n = 3 начинается так:

1, 3, 4, 9, 10, 12, 13, ...

Пятое число — это 10.

Как посчитать эффективно

Перебирать все допустимые числа нельзя: их слишком много.

Но нам это и не нужно.

Достаточно:

• идти по битам числа k;
• хранить текущую степень числа n;
• если текущий бит равен 1, прибавлять эту степень к ответу.

Что будем хранить

• ans — ответ;
• pw — текущая степень числа n.

Сначала:

• ans = 0
• pw = 1, потому что n^0 = 1

Дальше пока k > 0:

• если последний бит числа k равен 1, прибавляем pw к ответу;
• умножаем pw на n, чтобы перейти к следующей степени;
• сдвигаем k вправо на один бит.

Все вычисления делаем по модулю 10^9 + 7.

Алгоритм

Для каждого теста:

1. Считать n и k.
2. Инициализировать ans = 0, pw = 1.

3. Пока k > 0:

   • если k нечётно, то добавить pw к ans;
   • обновить pw = pw * n по модулю;
   • разделить k на 2.
4. Вывести ans.

Оценка сложности

У числа k не больше около 30 бит, потому что k <= 10^9.

Значит:

• время на один тест: O(log k)
• память: O(1)

Это очень быстро.

Эталонное решение на C++

#include <bits/stdc++.h>
using namespace std;

const long long MOD = 1000000007LL;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int t;
    cin >> t;

    while (t--) {
        long long n, k;
        cin >> n >> k;

        long long ans = 0;
        long long pw = 1;

        while (k > 0) {
            if (k & 1) {
                ans = (ans + pw) % MOD;
            }
            pw = (pw * n) % MOD;
            k >>= 1;
        }

        cout << ans << '\n';
    }

    return 0;
}

Эталонное решение на Python

MOD = 10**9 + 7

t = int(input())

for _ in range(t):
    n, k = map(int, input().split())

    ans = 0
    pw = 1

    while k > 0:
        if k & 1:
            ans = (ans + pw) % MOD
        pw = (pw * n) % MOD
        k >>= 1

    print(ans)

Что важно понять в этой задаче

Здесь ключевая мысль не в степенях сама по себе, а в связи между:

• двоичной записью числа;
• выбором степеней n.

Если вы замечаете, что каждый элемент либо берётся, либо не берётся, то очень часто это связано с битами.

В этой задаче бит числа k напрямую говорит, нужно ли добавлять степень n с таким номером.

Это и превращает задачу в короткое и красивое решение.

Итог

Чтобы найти k-е допустимое число:

• смотрим на двоичную запись k;
• для каждого бита 1 добавляем соответствующую степень n;
• считаем всё по модулю.

Это классическая задача на понимание двоичной записи числа и аккуратную реализацию