pythonでコーディング問題を解く3

問題

SNSのユーザー0-6と、他のユーザーとの距離を算出する関数を作成してください。

距離の定義について
A,Bが友達、B,Cが友達の場合
AとB、BとCの距離は1
AとCの距離は2
と定義する。

# 辞書配列で各ユーザーの友人関係を定義する
friends = {}
friends[0] = [1, 2]
friends[1] = [0]
friends[2] = [0, 3]
friends[3] = [2, 4]
friends[4] = [3, 6]
friends[5] = [6]
friends[6] = [4, 5]


def checkDistance(origin) :
    # 作業用キュー
    queue = []
    # すでに訪問したユーザーの距離
    distances = {}

    # 起点を格納する
    distances[origin] = 0
    queue.append(origin)

    while len(queue) != 0 :
        user = queue.pop()

        # friends配列にない場合スキップ              
        if not friends.get(user):
            continue

        # 距離が格納されていない、つまり未訪問の場合のみ距離を更新し、隣接する友人をqueueに追加する
        for friend in friends[user]:
            if distances.get(friend) is None :
                distances[friend] = distances.get(user) + 1
                queue.append(friend)

    return distances

print(checkDistance(0))
print(checkDistance(1))
print(checkDistance(2))
print(checkDistance(3))
print(checkDistance(4))
print(checkDistance(5))
print(checkDistance(6))