Strong friendship

# Table: Friendship
# 
# +-------------+------+
# | Column Name | Type |
# +-------------+------+
# | user1_id    | int  |
# | user2_id    | int  |
# +-------------+------+
# (user1_id, user2_id) is the primary key (combination of columns with unique values) for this table.
# Each row of this table indicates that the users user1_id and user2_id are friends.
# Note that user1_id < user2_id.
# 
#  
# 
# A friendship between a pair of friends x and y is strong if x and y have at least three common friends.
# 
# Write a solution to find all the strong friendships.
# 
# Note that the result table should not contain duplicates with user1_id < user2_id.
# 
# Return the result table in any order.
# 
# The result format is in the following example.
#
# Example 1:
# Input:
# Friendship table:
# +----------+----------+
# | user1_id | user2_id |
# +----------+----------+
# | 1        | 2        |
# | 1        | 3        |
# | 2        | 3        |
# | 1        | 4        |
# | 2        | 4        |
# | 1        | 5        |
# | 2        | 5        |
# | 1        | 7        |
# | 3        | 7        |
# | 1        | 6        |
# | 3        | 6        |
# | 2        | 6        |
# +----------+----------+
# Output:
# +----------+----------+---------------+
# | user1_id | user2_id | common_friend |
# +----------+----------+---------------+
# | 1        | 2        | 4             |
# | 1        | 3        | 3             |
# +----------+----------+---------------+
# Explanation:
# Users 1 and 2 have 4 common friends (3, 4, 5, and 6).
# Users 1 and 3 have 3 common friends (2, 6, and 7).
# We did not include the friendship of users 2 and 3 because they only have two common friends (1 and 6).

import pandas as pd

def strong_friendship(friendship: pd.DataFrame) -> pd.DataFrame:
  # Build bidirectional friendship
  bi = pd.concat([
    friendship[['user1_id', 'user2_id']],
    friendship[['user2_id', 'user1_id']].rename(columns={'user2_id': 'user1_id', 'user1_id': 'user2_id'})
  ])
  # For each pair in friendship, find common friends
  # A common friend c of (a, b) means (a, c) and (b, c) both exist in bi
  merged = friendship.merge(bi, left_on='user1_id', right_on='user1_id')
  # merged has user1_id, user2_id_x (original pair), user2_id_y (friend of user1_id)
  # Check if user2_id_y is also a friend of user2_id_x
  merged = merged.merge(bi, left_on=['user2_id_x', 'user2_id_y'], right_on=['user1_id', 'user2_id'])
  counts = merged.groupby(['user1_id_x', 'user2_id_x']).size().reset_index(name='common_friend')
  counts = counts[counts['common_friend'] >= 3]
  counts = counts.rename(columns={'user1_id_x': 'user1_id', 'user2_id_x': 'user2_id'})
  return counts[['user1_id', 'user2_id', 'common_friend']]