def _infection_indicator(
    health_states, infected_state_index):
  """Returns a binary vector that indicates whether individuals are infected."""
  return [int(state == infected_state_index) for state in health_states]