Задача с олимпиады

Необходимо написать программу, которая находит k-ое число по возрастанию в массиве A = <a1,a2,a3....an>. Массив А задаётся с помощью полинома P(x) = введите сюда описание изображения, ai = P(i) mod 2021.

Задано k и n.

Пример:

k = 1, n = 1, то ответ 169. k = 1, n = 10, то ответ 36.

Возможно ли оптимизировать вычисления, если да то как (варианта кроме как просто считать каждое число)? (ограничения k и n 4000000, время 1 секунда, память 256Мб). Возможно, следует использовать multiset, это единственное, что приходит в голову. Проблема переполнения типов решается преобразованием уравнения:

Из 123x^3 +456x^2+789*x+822%2021 в введите сюда описание изображения после каждой операции делать взятие остатка по модулю 2021.


Ответы (1 шт):

Автор решения: Stanislav Volodarskiy

Многочлен считаем по схеме Горнера. После каждого шага берём по модулю чтобы не было переполнения.

Про многочлен известно что p(x + 2021) = p(x) mod 2021. Поэтому считать его надо только 2021 раз.

counts[y] хранит количество раз когда y = p(x), 1 <= x <= n.

Функция in_range(a, b, x) считает сколько раз арифметическая прогрессия x + k * modulo попадает в полуинтервал [a, b).

Для каждой прогрессии (их 2021 различных) многочлен считаем один раз, получаем y. counts[y] увеличиваем на число элементов прогрессии в интервале [1, n].

Когда counts готов нужно последовательно вычитать его значения из k. Индекс на котором k станет не положительным - k-тое значение в отсортированной последовательности a_i, 1 <= i <= n.

#include <iostream>

const int modulo = 2021;

int p(int x) {
    int y = 123;
    y = (456 + x * y) % modulo;
    y = (789 + x * y) % modulo;
    y = (822 + x * y) % modulo;
    return y;
}

// number of {x + k * module} in half interval [a, b)
int in_range(int a, int b, int x) {
    return (b - x + modulo - 1) / modulo - (a - x + modulo - 1) / modulo;
}

int main() {
    int k;
    int n;
    if(!(std::cin >> k >> n)) {
        return 1;
    }

    int counts[modulo] = {0};
    for (int i = 0; i < modulo; ++i) {
        counts[p(i)] += in_range(1, n + 1, i);
    }

    for (int j = 0; j < modulo; ++j) {
        k -= counts[j];
        if (k <= 0) {
            std::cout << j << '\n';
            break;
        }
    }
}
$ g++ -std=c++11 -pedantic -Wall -Wextra -Werror fast.cpp

echo 1 1 | ./a.out
169

$ echo 1 10 | ./a.out
36

$ echo 1 4000000 | ./a.out
5

$ time echo 4000000 4000000 | ./a.out
2019

real  0m0.002s
user  0m0.000s
sys   0m0.000s
→ Ссылка