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

問題

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

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

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

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

[制約]
🔹$ 2 \leqq N \leqq 100000 $
🔹$ 1 \leqq A_i \leqq i-1 (2 \leqq i \leqq N) $

解き方・ソースコード

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

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

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

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

具体的には、社員 $i$ の持つ部下の社員番号が $j_1, j_2, \ldots, j_k$ である場合、$G[i]$ には $[j_1, j_2, \ldots, j_k]$ というリストが格納されます。

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

具体的には、各社員の持つ部下の数を $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 人

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