G - AtCoder Office Editorial by spheniscine


Let’s define a threshold \(\theta \approx \sqrt M\), and define every person that has at least that many records as heavy, with the rest being light.

There can thus be at most \(\frac M \theta \approx \sqrt M\) heavy people. For each heavy person, let’s precalculate the result for every potential person they may be paired with in a query. This can be done in \(O(N + M)\) for each heavy person, using prefix sums and a single pass through the records, so the total time spent in this step is \(O((N+M) \sqrt M)\).

We also separate the records into lists for each person for later use.

We can thus answer every query that involves at least one heavy person in \(O(1)\) each. The remaining queries involve two light people. To answer them, we use the separated records; since each light person has less than \(\theta \approx \sqrt M\) records, we can iterate through the relevant records for the two people involved and calculate the answer in \(O(\sqrt M)\) time.

Final complexity: \(O((N+M+Q) \sqrt M)\) time, \(O(N \sqrt M + M)\) space

Sample Code (Rust)

posted:
last update: