深さ優先探索（社員の階級）

問題

ある会社には、$N(\leqq 100000)$ 人の社員がいて、$1$ から $N$ までの番号が付けれらています。

この会社について、次の $N-1$ 個の情報が分かっています。

$i$ 個目の情報：社員 $ A_i$ と社員 $B_i$ は直属の上司と部下の関係にあり、社員 $ A_i$ と社員 $B_i$ のどちらが上司なのかは分かっていません。

社員 $T$ が社長であり、それ以外の $N-1$ 人全員が誰か1人の直属の部下であるとき各社員の階級を求めて下さい。

ただし、部下がいない社員の階級は $0$ であり、部下がいる社員の階級は、直属の部下における階級の最大値に $1$ を足した値であるとします。

解き方・ソースコード

前回記事の問題とは違い、今回の問題は社員番号の制約がないので、深さ優先探索を使って解いていきます。

まず入力値として、変数 $N$ には社員数、変数 $T$ には社長の社員番号、変数 $R$ には上司または部下の関係を設定します。

次に、階層構造を持つデータを扱うため、グラフの形式で社員の関係を表現します。

変数 $g$ には社員の関係を格納するために、空のリストで初期化し、上司・部下関係の情報をもとに関係を追加していきます。

その後、階級を格納するために配列 $rank$ を作成します。初期値はすべて0です。

関数dfsは、深さ優先探索アルゴリズムを用いて、各社員の階級を計算します。

引数 parentには、ある社員の上司の社員番号を設定します。

引数 iには、チェックする社員の社員番号を設定します。

この関数内では $dfs(i, j) + 1$により、部下の階級から自分の階級を計算しています。

そして、現在保持している階級よりも大きければ更新します。

最初の関数dfsの呼び出しは、社長の階級を求めるために $dfs(-1, T - 1)$とします。

最後に、すべての $rank$ の値を出力することで、各社員の階級を表示します。

[Google Colaboratory]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
#-------- 入力例１ ---------
N = 7       # 社員数
T = 3       # 社長の社員番号
R = [(3, 2), (3, 1), (2, 5), (1, 4), (4, 6), (4, 7)]  # 上司または部下の関係 
#-------- 入力例２ ---------
# N = 7       # 社員数
# T = 7       # 社長の社員番号
# R = [(3, 2), (3, 1), (2, 5), (1, 4), (4, 6), (4, 7)]  # 上司または部下の関係 
#---------------------------
g = [[] for i in range(N)]
for A, B in R:
    A -= 1
    B -= 1
    g[A].append(B)
    g[B].append(A)

rank = [0] * N

# 深さ優先探索
#   parent:上司の社員番号
#   i:チェックする社員番号
def dfs(parent, i):
    for j in g[i]:          # ある社員の部下全員分をループ
        if j == parent:     # 上司の社員番号は無視
            continue
        r = dfs(i, j) + 1   # チェックする社員番号を上司として、その部下のランクを取得する
        if rank[i] < r:     # 現在保持しているものよりランクが大きければ更新する
            rank[i] = r
    return rank[i]
dfs(-1, T - 1)  # 上司はいないので-1、社長番号T（インデックス化するため-1）としてランクを取得する

for i in range(N):
    print(f'社員 {i + 1} の階級は {rank[i]}')

[実行結果（入力例１）]

社員 1 の階級は 2

社員 2 の階級は 1

社員 3 の階級は 3

社員 4 の階級は 1

社員 5 の階級は 0

社員 6 の階級は 0

社員 7 の階級は 0

[実行結果（入力例２）]

社員 1 の階級は 3

社員 2 の階級は 1

社員 3 の階級は 2

社員 4 の階級は 4

社員 5 の階級は 0

社員 6 の階級は 0

社員 7 の階級は 5

各社員の階級を表示することができました。