nodes: 0, module; 1, function_definition; 2, function_name:relate_obs_ids_to_chosen_alts; 3, parameters; 4, block; 5, identifier:obs_id_array; 6, identifier:alt_id_array; 7, identifier:choice_array; 8, expression_statement; 9, comment:# Figure out which units of observation chose each alternative.; 10, expression_statement; 11, for_statement; 12, comment:# Return the desired dictionary.; 13, return_statement; 14, comment:"""
      Creates a dictionary that relates each unique alternative id to the set of
      observations ids that chose the given alternative.
  
      Parameters
      ----------
      obs_id_array : 1D ndarray of ints.
          Should be a long-format array of observation ids. Each element should
          correspond to the unique id of the unit of observation that corresponds
          to the given row of the long-format data. Note that each unit of
          observation may have more than one associated choice situation.
      alt_id_array : 1D ndarray of ints.
          Should be a long-format array of alternative ids. Each element should
          denote the unique id of the alternative that corresponds to the given
          row of the long format data.
      choice_array : 1D ndarray of ints.
          Each element should be either a one or a zero, indicating whether the
          alternative on the given row of the long format data was chosen or not.
  
      Returns
      -------
      chosen_alts_to_obs_ids : dict.
          Each key will be a unique value from `alt_id_array`. Each key's value
          will be a 1D ndarray that contains the sorted, unique observation ids
          of those observational units that chose the given alternative.
      """; 15, assignment; 16, identifier:alt_id; 17, call; 18, comment:# Determine which observations chose the current alternative.; 19, block; 20, identifier:chosen_alts_to_obs_ids; 21, identifier:chosen_alts_to_obs_ids; 22, dictionary; 23, attribute; 24, argument_list; 25, expression_statement; 26, comment:# Store the sorted, unique ids that chose the current alternative.; 27, expression_statement; 28, identifier:np; 29, identifier:sort; 30, call; 31, assignment; 32, assignment; 33, attribute; 34, argument_list; 35, identifier:selection_condition; 36, line_continuation:\; 37, call; 38, subscript; 39, line_continuation:\; 40, call; 41, identifier:np; 42, identifier:unique; 43, identifier:alt_id_array; 44, attribute; 45, argument_list; 46, identifier:chosen_alts_to_obs_ids; 47, identifier:alt_id; 48, attribute; 49, argument_list; 50, identifier:np; 51, identifier:where; 52, binary_operator:(alt_id_array == alt_id) & (choice_array == 1); 53, identifier:np; 54, identifier:sort; 55, call; 56, parenthesized_expression; 57, parenthesized_expression; 58, attribute; 59, argument_list; 60, comparison_operator:alt_id_array == alt_id; 61, comparison_operator:choice_array == 1; 62, identifier:np; 63, identifier:unique; 64, subscript; 65, identifier:alt_id_array; 66, identifier:alt_id; 67, identifier:choice_array; 68, integer:1; 69, identifier:obs_id_array; 70, identifier:selection_condition
code: def relate_obs_ids_to_chosen_alts(obs_id_array,
                                    alt_id_array,
                                    choice_array):
      """
      Creates a dictionary that relates each unique alternative id to the set of
      observations ids that chose the given alternative.
  
      Parameters
      ----------
      obs_id_array : 1D ndarray of ints.
          Should be a long-format array of observation ids. Each element should
          correspond to the unique id of the unit of observation that corresponds
          to the given row of the long-format data. Note that each unit of
          observation may have more than one associated choice situation.
      alt_id_array : 1D ndarray of ints.
          Should be a long-format array of alternative ids. Each element should
          denote the unique id of the alternative that corresponds to the given
          row of the long format data.
      choice_array : 1D ndarray of ints.
          Each element should be either a one or a zero, indicating whether the
          alternative on the given row of the long format data was chosen or not.
  
      Returns
      -------
      chosen_alts_to_obs_ids : dict.
          Each key will be a unique value from `alt_id_array`. Each key's value
          will be a 1D ndarray that contains the sorted, unique observation ids
          of those observational units that chose the given alternative.
      """
      # Figure out which units of observation chose each alternative.
      chosen_alts_to_obs_ids = {}
  
      for alt_id in np.sort(np.unique(alt_id_array)):
          # Determine which observations chose the current alternative.
          selection_condition =\
              np.where((alt_id_array == alt_id) & (choice_array == 1))
  
          # Store the sorted, unique ids that chose the current alternative.
          chosen_alts_to_obs_ids[alt_id] =\
              np.sort(np.unique(obs_id_array[selection_condition]))
  
      # Return the desired dictionary.
      return chosen_alts_to_obs_ids