UQR - Editorial

PROBLEM LINK:

Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4

Author: helloLad
Tester: sky_nik
Editorialist: iceknight1093

DIFFICULTY:

TBD

PREREQUISITES:

Math

PROBLEM:

You are given N, A, B. It is guaranteed that |A-B| \leq 1.
Find the number of integers X such that 1 \leq X \leq N, and there exist non-negative integers p and q such that Ap + Bq = X.

EXPLANATION:

Without loss of generality, let A \leq B.
Since |A-B| \leq 1, we must have either A = B or A+1 = B.

If A = B, the only numbers that can be formed are multiples of A, that is, A, 2A, 3A, \ldots
So, the answer is simply the number of positive multiples of A that don’t exceed A, which is \displaystyle \left\lfloor \frac{N}{A} \right\rfloor.

Now, we look at the A+1 = B case; in simpler words we want to know which numbers \leq N can be formed by adding copies of A and (A+1).
If you work out a few small cases, you might observe that:

  • It’s impossible to reach any X \lt A.
  • A and A+1 can be reached.
  • It’s impossible to reach any A+1 \lt X \lt 2A.
  • 2A, 2A+1, 2A+2 can be reached.
  • It’s impossible to reach any 2A+2 \lt X \lt 3A.
  • Everything from 3A to 3A+3 can be reached.
    \vdots

Extrapolating this, it seems useful to look at integers of the form pA + q.
Note that we only need to consider q \lt A, since if q \geq A we could just increase p by 1 and decrease q by A.
Henceforth, we only consider q \lt A.

From what was observed above, a reasonable guess is that we can reach pA + q if and only if p \geq q.
It’s not hard to prove the correctness of this.

Proof

If p\geq q, it’s easy to reach pA + q - take q copies of (A+1) and (p-q) copies of A.

if p \lt q, then combined with the fact that q \lt A, observe that we need at least q copies of A+1 — after all, copies of A don’t change the remainder modulo A while copies of (A+1) increase by 1, and we want the remainder modulo A to be q.
But then taking q copies of (A+1) already puts us at a value of qA + q, which is strictly larger than pA + q.
So, it’s impossible to form pA + q when p \lt q \lt A.

So, we want to count the number of integers \leq N that can be written as pA + q, for p \geq q and q \lt A.
(Note that for a fixed integer X, this choice of p and q is unique since we enforced 0 \leq q \lt A - they’re just the quotient and remainder when X is divided by A.)

That can be done quickly with a bit of math:

  • There are two valid values \lt 2A.
  • There are three valid values \lt 3A but \geq 2A.
  • There are four valid values \lt 4A but \geq 3A.
    \vdots

This continues till we reach A^2 = A\cdot A, where there are A valid values \lt A^2 and \geq A\cdot (A-1).
Then, for any m\geq A, there will always be only A (and not m) valid values between (m-1)A and mA-1, because there are only A values in that range at all.

So, let k = \left\lfloor \frac{N}{A} \right\rfloor.
Then,

  • If k \leq A, the count is 2 + 3 + \ldots + k.
  • If k \gt A, the count is 2 + 3 + \ldots + A + (k-A)\cdot A.

Both of these can be computed in constant time.
Note that these aren’t quite the final answers: we only considered blocks of length A, while the final block might be of length \lt A (in particular, it will have length (N\bmod A)).
This final block can be taken care of specially: p is fixed for it, so just count the number of valid q, which shouldn’t be hard if you’ve come this far.

TIME COMPLEXITY:

\mathcal{O}(1) per testcase.

CODE:

Author's code (C++)
#include <bits/stdc++.h>
using namespace std;
#define IOS ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
#define ll long long

int main(){
    IOS
    int t;
    cin>>t;
    while(t--){
        long long n,a,b;
        cin>>n>>a>>b;
        assert(abs(a-b)<=1);
        if(a==b){
            cout<<n/a<<'\n';
            continue;
        }
        if(a>b){
            swap(a,b);
        }
        long long k=n/a;
        if(k==0){
            cout<<0<<'\n';
            continue;
        }
        long long res=n;
        res-=(a-max(1LL,a-k))*(max(1LL,a-k)+a-1)/2;
        if(k<a-1){
            res-=max(0LL,n-a*k-k);
        }
        cout<<res<<'\n';
    }
    return 0;
}
Tester's code (C++)
#include <bits/stdc++.h>
using namespace std;

int main() {
  cin.tie(0)->sync_with_stdio(0);

  int t; cin >> t; while (t--) {
    int64_t n, a, b;
    cin >> n >> a >> b;

    if (a == b) {
      cout << n / a << '\n';
      continue;
    }

    if (a > b) {
      swap(a, b);
    }
    assert(b == a + 1);

    if (n < a) {
      cout << 0 << '\n';
      continue;
    }

    // Chicken McNugget Theorem
    if (n >= (a - 1) * a) {
      cout << n - ((a - 1) * a) / 2 << '\n';
      continue;
    }

    const auto quo = n / a;
    const auto rem = n % a;
    cout << (quo * (quo + 1)) / 2 + min(quo, rem) << '\n';
  }
}
Editorialist's code (Python)
for _ in range(int(input())):
    n, a, b = map(int, input().split())
    if a == b:
        print(n//a)
        continue
    a = min(a, b)
    k = n//a
    if k == 0:
        print(0)
        continue
    
    ans = 0
    if k >= a:
        ans = a*(a+1)//2 + (k-a)*a - 1
    else:
        ans = k*(k+1)//2 - 1
    ans += min(n%a, k)+1 # Last block
    print(ans)
1 Like

Can the tester explain what the Chicken McNugget Theorem is ?

1 Like

I recommend reading about Chicken McNugget Theorem on AOPS. It tells you two things: the largest number X^* that cannot be expressed as A \cdot p + B \cdot q, and how many numbers 1 \le X \le X^* cannot be expressed in this form. Therefore, one can use it to answer cases with N \ge X^* quickly.

4 Likes