Number of connected cities in a range

Question

There are N cities given in an array A. There is also a bike that can travel atmost K units between the cities.

There are Q queries that we need to answer. Each query has the form L R X. It asks for the number of cities that are reachable from the city X that are between L and R in A (1 - indexed). Each city has a petrol pump so you can assume that fuel gets replenished as you reach it.

Example:

A = [4, 3, 1, 9, 6], K = 2

Q1 = 1 3 6 => (3)

Q2 = 1 5 3 => (4)

In Q1, from city 6 you can go to city 4, then to city 3, and then to 1.

In Q2, from city 3 you can go to all of the cities except the city at 9.

Constraints:

N <= 10^5 and Q <= 10^5 and K <= 10^8

How do I solve this problem? Obviously it is not possible to do a DFS / BFS from each X because it is very costly and will timeout. I tried thinking of Disjoint sets to join the cities that are within K distance from each other but I don't have a very clear idea of it.

Any help is appreciated. Thank you!

Answer 1

I recommend:

1. Sort the cities by their location: [4, 3, 1, 9, 6] -> [1,3,4,6,9]
2. A linear scan through this will allow you to identify which sets of cities are reachable from each other. Label the cities according to which group they are in: [1,1,1,1,2]
3. Label each query according to the label of its starting city
4. Add all the label 1 cities into a segment tree
5. Use the segment tree to answer all the label 1 queries
6. Remove the label 1 cities from the segment tree

Then repeat steps 4,5,6 for each label.

This should give a total complexity of O(NlogN)+O(QlogN).