AtCoder が提供しているABC（AtCoder Beginner Contest）309 のE問題をC++とPythonで解いてみました。ABC309は、2023年7月8日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Family and Insurance（Difficulty : 957）
問題はリンク先をご覧ください。
ABC309 E問題 Family and Insurance
自分の親を調べて、自分の状況を更新すれば解くことができます。AtCoder Problems による Difficulty は 957 でした。
解答案
DFS で解く（結果は TLE）
最初に TLE（Time Limit Exceeded）判定となるプログラムを紹介します。基本的な考え方は以下です。
- DFS を使う。引数として頂点番号（X[i]）と補償対象世代数（y[i]）を渡す。
- 補償内であれば、result[i] を true にして、子供に対して再帰的に呼び出す。
- result[i］が true の個数が解答となる。
#include <bits/stdc++.h>
using namespace std;
vector<vector<int>> G;
vector<int> x;
vector<int> y;
vector<bool> result;
void dfs(int v, int dist)
{
if (dist >= 0) {
result[v] = true;
} else {
return;
}
for (auto next_v : G[v]) {
dfs(next_v, dist - 1);
}
}
int main()
{
int n, m;
cin >> n >> m;
G.resize(n + 1);
for (int i = 2; i <= n; ++i) {
int p;
cin >> p;
G[p].push_back(i);
}
x.resize(m);
y.resize(m);
for (int i = 0; i < m; ++i) {
cin >> x[i] >> y[i];
}
result.assign(n + 1, false);
for (int i = 0; i < m; ++i) {
dfs(x[i], y[i]);
}
int out_result = 0;
for (int i = 1; i <= n; ++i) {
if (result[i]) {
++out_result;
}
}
cout << out_result << endl;
return 0;
}残念ながら、いくつかのテストケースでTLEとなりました。y[i] を大きい順から処理するように書き直しましたが定数倍しか改善できず結果は TLE でした。
問題の制約を振り返る
以下の制約がポイントになります。
$$1 \leqq p_i \leqq i – 1$$
つまり、自分の親の頂点は、番号が自分より少なくなります。自分自身が補償対象になっているかは、自分の親と自分自身について調べればよいため、DP（動的計画法）に近い解き方が適用できます。
変数 dp[i] は、補償対象世代数とします。
- 問題の入力から自分自身の dp[i] を更新する。
- 頂点1からNまで、自分の親と自分自身の値から dp[i] を更新する。このとき、i を小さい方から処理すればよい。
C++ プログラム例（ABC309E）
上記の考察をプログラムに落とし込みます。親を表す配列 p の添え字は2からNまでは問題で与えられます。p[1]=1 としました（9行目）。
DFS の解法よりもすっきりと書けています。変数 dp に関係する行の背景色を変更しておきました。
以下が、C＋＋プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
int main()
{
int n, m;
cin >> n >> m;
vector<int> p(n + 1, 0);
p[1] = 1;
for (int i = 2; i <= n; ++i) {
cin >> p[i];
}
vector<int> dp(n + 1, -1);
for (int i = 0; i < m; ++i) {
int x, y;
cin >> x >> y;
dp[x] = max(dp[x], y);
}
for (int i = 1; i <= n; ++i) {
dp[i] = max(dp[i], dp[p[i]] - 1);
}
int result = 0;
for (int i = 1; i <= n; ++i) {
if (dp[i] >= 0) {
++result;
}
}
cout << result << endl;
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC309E）
Python も C++ と同じ解法です。リスト p は、添え字を問題に合わせるために 0 番目と 1 番目の要素を挿入しています（4、5行目）。
"""AtCoder Beginner Contest 309 E"""
n, m = map(int, input().split())
p = list(map(int, input().split()))
p.insert(0, 0)
p.insert(1, 1)
dp = [-1] * (n + 1)
for i in range(m):
x, y = map(int, input().split())
dp[x] = max(dp[x], y)
for i in range(1, n + 1):
dp[i] = max(dp[i], dp[p[i]] - 1)
result = 0
for i in range(1, n + 1):
if dp[i] >= 0:
result += 1
print(result)こちらも「AC」と判定されました。
最後に
今回のコンテストは、開始44分でD問題までACを取ることができました。残り時間でE問題に取り組みましたが、TLEを解消することができませんでした。
公式解説から、自分の親と自分の状況から、自分の状況が更新できることが理解できました。グラフの問題でしたが、DPに近い解法でした。学ぶことができました。
引き続き ABC の問題を紹介していきます。
