POI Task Distance

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

/* This version of the code is adapted with modified precomputation
 - only precomputing up to necessary max for specific testcase
 This is because there are some small tests with low time limits */

int main() {
    ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);

    int n, mx = 0; cin >> n;
    vector<int> a(n);
    for (int &x : a) cin >> x, mx = max(mx, x);

    vector<int> smallestFactor(1+mx); // smallest prime factor
    vector<vector<int>> pos(1+mx); // list of positions in the array of each value up to MX
    vector<bool> got(1+mx);

    // Sieve of Eratosthenes
    for (int i = 2; i <= mx; ++i) {
        if (smallestFactor[i]) continue;
        smallestFactor[i] = i;
        for (long long j = (long long)i*i; j <= mx; j += i) {
            if (!smallestFactor[j] || smallestFactor[j] > i) smallestFactor[j] = i;
        }
    }

    auto getF = [&smallestFactor](int x) { // sum of powers of prime factors
        int f = 0;
        while (smallestFactor[x]) {
            x /= smallestFactor[x];
            ++f;
        }
        return f;
    };

    vector<int> f(n);
    for (int i = 0; i < n; ++i) {
        pos[a[i]].push_back(i);
        got[a[i]] = true;
        f[i] = getF(a[i]);
    }

    vector<int> ans(n, -1), ansDist(n, -1);

    auto lowerF = [&f](int i, int j) {
        return j == -1 || f[i] < f[j] || (f[i] == f[j] && i < j);
    };

    // go through all divisors (will be gcd in each final answer)
    for (int d = 1; d <= mx; ++d) {
        int fd = getF(d);

        // consider all multiples of d in the array
        int small = -1, next = -1; // smallest, next smallest
        for (int m = d; m <= mx; m += d) {
            if (!got[m]) continue;
            for (int i : pos[m]) {
                if (lowerF(i, next)) {
                    next = i;
                    if (lowerF(next, small)) swap(small, next);
                }
            }
        }

        for (int m = d; m <= mx; m += d) {
            if (!got[m]) continue;
            for (int i : pos[m]) {
                int cur = small;
                if (i == small) {
                    if (next == -1) continue;
                    cur = next;
                }
                int dist = f[i] + f[cur] - 2*fd;

                if (ans[i] == -1 || dist < ansDist[i] || (dist == ansDist[i] && cur < ans[i])) ans[i] = cur, ansDist[i] = dist;
            }
        }
    }

    for (int i = 0; i < n; ++i) cout << ans[i]+1 << endl;
    return 0;
}