Example OpenAI generation

example-openai-output.py

def compute_from_graph(nodes, edges):
  """the closest female friends for each man"""
  nodes = nodes.set_index('id')
  # Extract the coordinates from the location column.
  nodes['latitude'] = nodes['location'].apply(lambda x: x[0])
  nodes['longitude'] = nodes['location'].apply(lambda x: x[1])
  # Compute the distance between each pair of nodes.
  distances = pd.DataFrame(np.zeros((len(nodes), len(nodes))), index=nodes.index, columns=nodes.index)
  for i, row_i in nodes.iterrows():
    for j, row_j in nodes.iterrows():
      if i < j:
        distances.loc[i, j] = np.sqrt((row_i['latitude'] - row_j['latitude']) ** 2 + (row_i['longitude'] - row_j['longitude']) ** 2)
        distances.loc[j, i] = distances.loc[i, j]
  # Keep only the distances between men and women.
  men = nodes[nodes['gender'] == 'Male']
  women = nodes[nodes['gender'] == 'Female']
  distances = distances.loc[men.index, women.index]
  # Find the closest female friend for each man.
  closest_women = pd.DataFrame({'distance': distances.min(axis=1), 'closest_woman': distances.idxmin(axis=1)})
  closest_women = closest_women.reset_index().rename(columns={'id': 'src'})
  edges['src'] = edges['src'].astype(str)
  edges['dst'] = edges['dst'].astype(str)
  edges = edges.merge(men.reset_index()[['id', 'name']], left_on='src', right_on='id')
  edges = edges.merge(closest_women, left_on='src', right_on='src')
  edges = edges.merge(women.reset_index()[['id', 'name']], left_on='closest_woman', right_on='id')
  return edges[['name_x', 'name_y']].rename(columns={'name_x': 'name', 'name_y': 'friend'})