data_science_from_scratch/ch1/friends.py

from typing import List, Dict, Tuple

USERS = [
    {"id": 0, "name": "Hero"},
    {"id": 1, "name": "Dunn"},
    {"id": 2, "name": "Sue"},
    {"id": 3, "name": "Chi"},
    {"id": 4, "name": "Thor"},
    {"id": 5, "name": "Clive"},
    {"id": 6, "name": "Hicks"},
    {"id": 7, "name": "Devin"},
    {"id": 8, "name": "Kate"},
    {"id": 9, "name": "Klein"},
]

FRIENDSHIPS = [
    (0, 1),
    (0, 2),
    (1, 2),
    (1, 3),
    (2, 3),
    (3, 4),
    (4, 5),
    (5, 6),
    (5, 7),
    (6, 8),
    (7, 8),
    (8, 9),
]


def add_friendships(
    users: List[Dict], friendships: List[Tuple[int, int]]
) -> List[Dict]:
    for user in users:
        user["friends"] = []

    for i, j in friendships:
        users[i]["friends"].append(users[j])  # add j as friend of i
        users[j]["friends"].append(users[i])  # add j as friend of i

    return users


def number_of_friends(user: Dict) -> int:
    return len(user["friends"])


def friend_counts(users: List[Dict]) -> List[Tuple[int, int]]:
    """Returns a sorted list by number of friends"""
    number_of_friends_by_id = [(user["id"], number_of_friends(user)) for user in users]

    _sorted_list = sorted(number_of_friends_by_id, key=lambda tup: tup[1], reverse=True)
    return _sorted_list


if __name__ == "__main__":
    users = add_friendships(USERS, FRIENDSHIPS)
    total_connections = sum(number_of_friends(user) for user in users)
    num_users = len(users)
    avg_connections = total_connections / num_users
    counts = friend_counts(users)

    print(f"Total user count: {num_users}")
    print(f"Total number of connections: {total_connections}")
    print(f"Average connections: {avg_connections}")
    print(counts)