covid_pop_data_structure.py

from bisect import bisect_left, bisect_right


class COVIDPopulation:
    def __init__(self, N, K):
        """
            K is the number of infected persons in N number of people.
            It's not a percentage.
        """
        self.N = N
        
        # Now we only need to store the indices of people who are infected.
        # So we generate those via np.random.choice
        self._people_infected = np.sort(np.random.choice(N, K, replace=False))
    
    def is_infected(self, left: int, right: int):
        """
            Check if there are numbers in self._people_infected that are between
            left and right, inclusive.
        """
        left_dex = bisect_left(self._people_infected, left)
        right_dex = bisect_right(self._people_infected, right)
        
        return (right_dex-left_dex) > 0

    def __str__(self):
        return "Infected: " + str(self._people_infected) + f" out of {self.N}"

    def __repr__(self):
        return self.__str__()