友達の友達

こんにちは。CPPXのCです。

今日も今日とてC問題です。
今回は特につまることもなくすんなりと解くことができたのですが、人間は日々退化していくものなので解き方を忘れないように書きとどめておくことにします。
自分のかいたソースコード、あとから見てもなんのことやらわからないんですよね。不思議と。

概要

さて、今日の問題はこちらです。

C: 友達の友達 - AtCoder Beginner Contest 016 | AtCoder

掻い摘んで説明すると、まずN人の人間のIDと、その友達の組がM組入力として与えられます。このとき、あるIDの人間の友達の友達は何人いるかをすべて出力しなさい、というのがこの問題です。

解法

Nが10人以下なので全探索で…とも思いましたが、友達の友達を調べるためにはただ一つ先のノードを見ればいいだけなので、全探索を使うまでもなく解けそうですね。
ということで2重で回して実装してみることにしました。

まずは入力を受け取ります。自分のIDの位置に友達すべてのIDのリストをつっこんだリストを生み出します。こうすることでIDに対応する友達のリストが生まれるはずです。

N, M = map(int, input().split())

input_list = []

for i in range(M):
    input_list.append(list(map(lambda x: int(x) - 1, input().split())))

data_list = [[] for i in range(N)]

for i in range(M):
    data_list[input_list[i][0]].append(input_list[i][1])
    data_list[input_list[i][1]].append(input_list[i][0])

こうして生まれたリストから友達の友達を調べます。最終的には友達の友達の数を調べたいので、
自分のID➡自分の友達のIDのリスト➡友達の友達のIDのリスト
という感じで中身を持ってきます。2重ループをつくってリストの中身で回せばよさそうですね。

result = []

for i, v in enumerate(data_list):
    tmp_list = []
    for j in v:
        tmp_list = list((set(data_list[j]) | set(tmp_list)) - set(v))
        tmp_list.remove(i)
    result.append(tmp_list)

for i in result:
    print(len(i))

最終的にほしいのはすべてのIDに対応する友達の友達の数なので、友達の友達のリストを格納するリストを作ってこのリストで回して出力することにしました。
ほしいのは友達の友達の数、すなわちリストの長さなので、このリストの中身の長さを順番に出力すればよいというわけです。

肝心の友達の友達リストの取得方法ですが、単純に自分の友達の友達をすべて論理和によって和集合として取得し、そのうち自身と、自身の友達を集合から抜くことで作ることができました。
Pythonだと便利な機能がいっぱいあって簡単に解けますね。 これ、Javaだとどう解くんだろう…