TREATS - Editorial


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

Author: raysh07
Tester: airths
Editorialist: iceknight1093






You have N boxes each of chocolates and candies. The counts of candies in each of these boxes are represented by arrays A and B.
Count the number of ways to choose exactly one box of candies and one box of chocolates so that the total number of treats can be distributed equally among M children.


If candy box i and chocolate box j are chosen, Chef will have a total of (A_i + B_j) chocolates.
This can be distributed equally to M children if and only if it’s a multiple of M.

Let’s fix i, the candy box that’s chosen.
Now, we’d like to count the number of j such that (A_i + B_j) is a multiple of M.
In other words, (A_i + B_j) should have a remainder of 0 when divided by M.

Observe that for this to be the case, B_j should have a remainder that’s “opposite” to that of A_i.
That is, the remainder when B_j is divided by M should be (M - A_i)\pmod M.
It’s not hard to see that we can choose any B_j that satisfies this property.

Let C be a 0-indexed array of length M, where C_x counts the number of indices j such that
B_j\pmod M = x

Then, following from our initial observations, for each i from 1 to N the number of valid indices j is exactly C_{(M - A_i)\pmod M}.

Computing the array C can be done in linear time, and finding the answer after that also takes linear time since we use a single lookup to array C for each i.


\mathcal{O}(N + M) per testcase.


Editorialist's code (Python)
for _ in range(int(input())):
    n, m = map(int, input().split())
    a = list(map(int, input().split()))
    b = list(map(int, input().split()))
    ct = [0]*m
    for i in range(n):
        a[i] %= m
        b[i] %= m
        ct[a[i]] += 1
    ans = 0
    for i in range(n):
        ans += ct[(m - b[i])%m]

what if we store mod of both choc boxes and candy boxes in a map…and then
we run a loop of 0 to m-1 cause every value corresponding to key will be less than m…and then we multiply the remainder of choc box to complement remainder(m-x) of candy box …and sum all up…keeping in mind the case of 0…i tried it but it was wrong on last test case…can u please help me with it…

It should work. The approach is fine. I did the same. Maybe there is some other tiny but in your implementation ?

i am sending u the code…would u please check it

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

int main() {

int t;
cin >> t;

while (t--) {
    int N, M;
    cin >> N >> M;
    vector<int> A(N), B(N);
    for (int i = 0; i < N; i++) {
        cin >> A[i];
    for (int i = 0; i < N; i++) {
        cin >> B[i];

    unordered_map<int, int> candyModCount, chocolateModCount;
    for (int i = 0; i < N; i++) {
        candyModCount[A[i] % M]++;
        chocolateModCount[B[i] % M]++;

    long long ans = 0;
    for (int i = 0; i < M; i++) {
        int complement = (M - i) % M;
        ans += candyModCount[i] * chocolateModCount[complement];

    cout << ans << endl;

return 0;


Yes it’s failing due to integers overflow on the line where you have done the multiplication.
You might think that ans variable is long long, so it shouldn’t overflow.
Reason this happens is that the evaluated temporary value of RHS is alloted int only by the compiler since both operands are ints.
So, before LHS (ans variable) gets updated, overflow has already happened in RHS. So the value being added to ans variable is itself incorrect.
You can try explicit type casting to long long in the line where you have done multiplication.
Fun experiment: try making one of your maps in long long data type and keep the other map as a simple int. This will pass because this time compiler does implicit type promotion while evaluating the temporary value in RHS and hence all goes fine.

To be honest, i stick to long long all the time.

Explicit type casting version
Different maps version

thanks buddy

1 Like