動的計画法（部下の人数）

問題

ある会社には N 人の社員がいて、地位順に 1 から N までの番号が付けられています。

社長（社員番号 1）以外には直属の上司が１人いて、社員 i の直属の上司は社員 Ai です。

各社員について、部下が何人いるかを出力してください。

ただし、社員 y が社員 x の部下であることは、x≠y であり、なおかつ社員 y の直属の上司をたどって社員 x に到達できることを意味します。

[制約]
🔹2≦N≦100000
🔹1≦Ai≦i−1(2≦i≦N)

解き方・ソースコード

動的計画法を使って、この問題を解いていきます。

まず、社員の数 N と直属の上司のリスト A を定義します。

次に、隣接リスト G を作成します。

隣接リストは、それぞれの社員が持つ部下の社員番号を管理します。

具体的には、社員 i の持つ部下の社員番号が j1,j2,…,jk である場合、G[i] には [j1,j2,…,jk] というリストが格納されます。

次に、動的計画法を用いて各社員が持つ部下の数を求めます。

具体的には、各社員の持つ部下の数を dp というリストに保存します。

最初は全員が部下を持たないものとして、全ての要素を 0 に初期化しています。

そして、社員 i の持つ部下の数を計算するために、G[i] の要素（つまり社員 i の部下の社員番号）を順番に取り出して、その社員が持つ部下の数を dp から参照して加算します。

最後に、dp[i] に 1 を加えることで、社員 i の直属の部下の数を dp に追加しています。

この処理を、社員 i が社長（社員番号１）である場合まで繰り返します。

最後に、dp の各要素を順番に出力することで、各社員が何人の部下を持つかを表示します。

[Google Colaboratory]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
#-------- 入力例１ ---------
N = 7                   # 社員数
A = [1, 1, 3, 2, 4, 5]  # 社員番号２から社員番号７の直属の上司の社員番号
#---------------------------
# 隣接リストの作成
G = [list() for i in range(N + 1)]
for i in range(2, N + 1):
    G[A[i - 2]].append(i) # 上司 → 部下の方向に方向に辺を追加

# 動的計画法（dp[x] は社員 x の部下の数）
dp = [0] * (N + 1)
for i in range(N, 0, -1):
    for j in G[i]:
        dp[i] += (dp[j] + 1)

for i in range(1, N + 1):
    print(f'社員 {i} の部下の数は {dp[i]} 人')

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

社員 1 の部下の数は 6 人

社員 2 の部下の数は 2 人

社員 3 の部下の数は 2 人

社員 4 の部下の数は 1 人

社員 5 の部下の数は 1 人

社員 6 の部下の数は 0 人

社員 7 の部下の数は 0 人

各社員について、部下が何人いるかを出力することができました。