Καλύτερη λύση στο happy (27ος camp seniors)

[doublemidi]

Βρήκα εναλλακτική και καλύτερη λύση για το πρόβλημα που πέρνει $Ο(1)$ ανά ερώτημα, αντί για το $Ο(sqrt(maxA))$ που έχει στον κώδικα του github ($maxA$ είναι η μέγιστη τιμή για $Α, Β$ στα ερωτήματα). Συγγνώμη αν είναι λάθος μέρος να το ανεβάσω αυτό (στα issues), δεν είμαι καθόλου εξικιωμένος με το github. Στο υπόλοιπο του post εξηγώ την λύση, μαζί με κώδικα.

Αρχικά, όταν μας δίνεται το ερώτημα $A$ $B$, μπορούμε να υπολογίσουμε το πλήθος των χαρούμενων στο $[1, Α-1]$ και στο $[1, Β]$ και να τους αφερέσουμε.
χ.β.γ., κοιτάμε το $[1, Α]$.

Θέτουμε $S(x)$, το άθροισμα των τετραγώνων των ψηφίων του $x$.
Θέτουμε $f(x)$, το "happiness" του $x$. Αν ο $x$, $f(x)=1$, αλλιώς $f(x)=0$

Παρατήρηση 1: $S(x) \leq 9*9*9=729$ για $x\leq maxA$ και $S(x) \leq 6*9*9=486$ για $x<=1'000'000$.
Κάνουμε precompute την $f(x)$, για $x \leq 729$. Δεν θα εξηγήσω πώς, γιατί αυτό γίνεται και στον κώδικα του pdp.

Τώρα, χωρίζουμε το εύρος $[1, 10^9]$ σε $1'000'000$ χιλιάδες.
Κάνουμε precompute την $S(x)$, για $x \leq 1'000'000$. Αυτό μπορεί να γίνει σε $Ο(n)$ (αντί για $O(nlogn)$) (εδώ $n=1'000'000$), με dp, όπου $dp()$ = $S()$, ως εξής:
$S(0)=0$
$S(x) = (x\%10)^2 + S(x/10)$, δηλαδή πέρνουμε κάποιο ψηφίο του $x$ (το τελευταίο σε αυτήν την περίπτωση), και αθροίζουμε το τετράγωνο του ψηφίου αυτού μαζί με την $S$ του $x$, χωρίς το ψηφίο αυτό.

Παρατήρηση 2: $ f(x) = f(S(x))$
Έτσι, μπορούμε να υπολογίσουμε την $f(x)$, για $x \leq 1'000'000$ σε $Ο(1)$.

Θέτουμε $g(i) = \sum\limits_{n=0}^{999} f(i*1000 + n)$, δηλαδή το άθροισμα της $g()$ για τιμές στην $i$-οστή χιλιάδα.
Θέλουμε να κάνουμε precompute το $g(i)$ για κάθε $0 \leq i \leq 999'999$.
Παρατήρηση 3: $S(i)=S(j) => g(i)=g(j)$.
Αυτό είναι προφανές, αφού $S(i)=S(j) =&gt; f(1000*i + x) = f(1000*j + x) = f(S(i)+S(x))$, για x<1000, αφού τα ψηφία των $i/j$ και του $x$ είναι "ανεξάρτητα".

Αρκεί λοιπόν να υπολογήσουμε μία τιμή της $g()$ για κάθε πιθανή τιμή της $S()$.
Θέτουμε λοιπόν $h(S(i)) = g(i)$, έτσι ώστε $S(i)=x$. Η $h()$ στον κώδικα λέγεται "thousands_to"
Υπολογίζουμε την $h(x)$, για $x\leq486$ brute force
Έπειτα υπολογίζουμε την $g(i), i&lt;1'000'000$, σε $Ο(1)$ για κάθε $i$ με την βοήθεια της $h()$.

Για κάθε ερώτημα, θέλουμε να υπολογίσουμε ένα prefix sum της $g()$, και απο την τελευταία χιλιάδα, κάποιο κομμάτι της. Για παράδειγμα, αν στο ερώτημα έχουμε $Α=3422$, η απάντηση είναι $g(0) + g(1) + g(2) +$ τις υπόλοιπες $422$ τιμές που θα δούμε σε λίγο πως θα τις υπολογίζουμε.
Για το prefix sum της $g()$, μπορούμε πολύ απλάνα αποθηκεύσουμε τα prefix sums της $g()$. Στον κώδικα, ο prefix sum πίνακας στον κώδικα λέγεται "thousands".

Θέτουμε $p(i, x) = \sum\limits_{n=0}^{999} f(i*1'000'000 + n)$
Παρόμοια με την παρατήρηση 3, έχουμε $S(i)=S(j) =&gt; p(i, x) = p(j, x)$.
Oμοίως, αρκεί να υπολογίσουμε μόνο $486$ τιμές της $p(i, x)$ και, ομοίως,
θέτουμε $q(S(i), x) = p(i, x)$.
Υπολογίζουμε την $q()$ για $i \leq 468$ brute force και ομοίως, αποθηκεύουμε το prefix sum (ως προς $x$) της $q$.
O prefix sum πίνακας στον κώδικα λέγεται "therest".
Για κάθε ερώτημα λοιπόν, όταν θέλουμε να υπολογίσουμε το prefix sum της $p(i, x)$, είναι ίσο με το prefix sum της $q(S(i), x)$.

Τέλος, παραλείποντας κάποιες λεπτομέριες, η απάντηση του ερωτήματος είναι thousands[x/1000] + therest[S[x/1000]][x%1000].
Οι λεπτομέριες αυτές είναι το αν οι πίνακες μας είναι $0$ ή $1$-indexed, και κάποια edge cases, όταν $A=maxA$ ή όταν $Α&lt;1'000$

Σημείωση: Πολλές φορές μπορεί να μην υπολογίζουμε την $f(0)$, αλλά $f(0)=0$, άρα δεν μας νοιάζει αν το υπολογίζουμε.

Κώδικας:

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

typedef long long ll;
typedef long double ld;
#define vc vector
typedef vc<ll> vcll;
typedef vc<bool> vcb;
#define pr pair
typedef pr<ll, ll> prll;
typedef map<ll,ll> mapll;
typedef unordered_map<ll,ll> umapll;

#define f0r(i,a,n) for ((i)=(a); (i)<=(n); (i)++)
#define r0f(i,n,a) for ((i)=(n); (i)>=(a); (i)--)

#define pb push_back
#define ppb pop_back
#define pf push_front
#define ppf pop_front
#define mp make_pair
#define fi first
#define se second
#define sz size
#define all(x) (x).begin(), (x).end()
#define all0(x, n) (x).begin(), (x).begin()+n
#define all1(x, n) (x).begin()+1, (x).begin()+n+1
#define allrev(x) (x).rbegin(), (x).rend()
#define in(v, s) ((s).find((v)) != (s).end())
#define GCD(x, y) __gcd(abs((x)), abs((y)))

#define INF (LLONG_MAX>>3ll)

#define maxO (9ll*9*9)
#define smallmaxO (9ll*9*6)

inline void maxa(ll &a, ll b) { if (a<b) a=b; }
inline void mina(ll &a, ll b) { if (a>b) a=b; }

ll q;
vcll S(1000010);
vcll f(maxO+10, -1);
vcb seen(maxO+10, 0);

vcll thousands(1000010); // psa
vcll thousands_to(smallmaxO+10);
ll therest[smallmaxO+10][1010]; // [add][pfx]

inline void calc_s()
{
	S[0]=0;
	ll i;
	f0r(i,1,1000000)
	{
		ll d=i%10;
		S[i] = S[i/10] + d*d;
	}
}

inline ll calc_f (ll x)
{
	if (f[x]!=-1) return f[x];
	if (seen[x]) return f[x]=0;
	if (x==1) return f[1]=1;
	seen[x]=1;
	f[x] = calc_f (S[x]);
	seen[x]=0;
	return f[x];
}
inline void calc_f()
{
	ll i;
	f0r(i,0,maxO) calc_f(i);
}

inline void build_to()
{
	ll s;
	f0r(s,0,smallmaxO)
	{
		ll cnt=0;
		ll a;

		f0r(a,0,999)
		{
			s+=S[a];
			cnt+=f[s];
			s-=S[a];
		}

		thousands_to[s]=cnt;
	}
}

inline void build_thousands()
{
	build_to();
	ll i;
	f0r(i,0,1000001)
	{
		ll cnt = thousands_to[S[i]];
		thousands[i] = (i==0 ? cnt : thousands[i-1] + cnt);
	}
}

inline void build_therest()
{
	ll s, i;

	f0r(s,0,smallmaxO) f0r(i,0,999)
	{
		bool v = f[S[i] + s];
		therest[s][i] = (i==0 ? v : therest[s][i-1] + v);
	}
}

inline ll qry (ll x) // psa
{
	ll i=x/1000, r=x%1000;
	i--;
	return (i==-1 ? 0 : thousands[i]) + therest[S[i+1]][r];
}

inline void queries()
{
	scanf("%lli", &q);
	while (q--)
	{
		ll l, r;
		scanf("%lli%lli", &l, &r);
		ll ans=0;
		if (r==1000000000) r--, ans=1;

		ans += qry(r) - qry(l-1);
		printf("%lli\n", ans);
	}
}

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

	freopen("happy.in", "r", stdin);
	freopen("happy.out", "w", stdout);

	calc_s();
	calc_f();
	build_thousands();
	build_therest();
	queries();

	return 0;
}

[Dim131]

Πολύ ωραία λύση και καλογραμμένη εξήγηση!

Υ. Γ. Υπάρχει ένα μικρό μπέρδεμα με το x και το j στο "Θέτουμε p(i, x) .."

Υ. Γ. 2. Μετέφερα το issue στο discussion (που πριν δεν ήταν ανοιχτό).
Ελπίζω και άλλοι να γράψουν για τις λύσεις τους είτε εδώ είτε στο φόρουμ.

[doublemidi]

Ευχαριστώ για την διόρθωση, το έφτιαξα