yutasの競技プログラミング勉強帖

競技プログラミングの問題についての解説記事を主に書いています。

C - Anji's Binary Tree / Codeforces Round 911 (Div. 2)

CF1900 Codeforces C++ Div.2 Codeforces Round 911 Diff 1300 木構造 DP 再帰 競技プログラミング

問題

頂点 1 を根とした n 頂点の二分木がある。

各頂点は最大で2個の子を持っており、また頂点 i は文字 {s}_{i} をそれぞれ持っている。 ただし、 {s}_{i} は U, L, R のいずれかである。

この二分木について、頂点 1 から以下のように動くことを考える。

  • いまいる頂点での文字が U ならば、その親の頂点に移動する。存在しない場合は何もしない。
  • いまいる頂点での文字が L ならば、その左側の子の頂点に移動する。存在しない場合は何もしない。
  • いまいる頂点での文字が R ならば、その右側の子の頂点に移動する。存在しない場合は何もしない。

この二分木に対して、以上の動作から葉に辿り着くためには、 {s}_{i} を何個変えなければならないか、その最小値を求めよ。 ただし、「葉」とは「子を持たない頂点」を指す。

入力

まず最初の1行目に、テストケースの個数を表す整数 t が与えられる。

その後、 t 個のテストケースのそれぞれに対して、以下のように入力が与えられる。

  • 最初の1行目には、整数 n が与えられる。
  • 次の1行には、 n 文字の文字列 s が与えられる。
  • 次の n 行に対して i 行目には、頂点 i の左側の子の頂点 {l}_{i}、右側の子の頂点 {r}_{i} がそれぞれ与えられる。
    • ただし、この値が 0 のときは、そのような頂点が存在しないということを表す。

条件

  • 1 \leq t \leq 5 \cdot {10}^{4}
  • 2 \leq n \leq 3 \cdot {10}^{5}
    • すべてのテストケースにおける n の値の合計は 3 \cdot {10}^{5} を超えない。
  • 0 \leq {l}_{i}, {r}_{i} \leq n

出力

各テストケースでの答えを1行ごとに出力すること。

解法

この問題では出発地点が頂点 1 であり、ゴールが「子を持たない頂点」と決まっています。 頂点 1 がどの頂点よりも親の位置にあたるため、「親を辿る」という操作は「葉に辿り着く」という目標に対して無駄な操作になってしまうと言えます。

ここで、 \mathrm{dp}_{i} を「頂点 i から出発し、子を持たない頂点に辿り着くまでに、文字列 s 内の文字を変える必要のある最小の個数」として定めます。

頂点 i の左側の子が存在する（すなわち {l}_{i} \ne 0）と仮定すると、頂点 i からこの左側の子を通って葉に辿り着くために変える必要のある個数は

  • {s}_{i} = L のとき: {s}_{i} は変える必要がないので、 \mathrm{dp}_{{l}_{i}}
  • {s}_{i} \ne L のとき: {s}_{i} を L に変える必要があるので、 \mathrm{dp}_{{l}_{i}} + 1

ということにそれぞれなります。 またこれは、右側の子が存在する場合についても同様のことが言えます。（ {s}_{i} = R かどうかで判断を行います。）

さらに、頂点 i が葉の場合については、最初からゴール地点に到達しているため、 \mathrm{dp}_{i} = 0 であると言えます。

従って、頂点 1 からそれぞれの頂点の子について二分木を順に見ていき、それぞれの頂点の \mathrm{dp} の値を更新することによって得られる \mathrm{dp}_{1} の値が答えとなります。

与えられた二分木について、頂点 1 から辺を経由して順番に見ていくことから、これは O(n) の計算量がかかりますが、 n \leq 3 \cdot {10}^{5} なので、実行時間制限に十分に間に合います。

ソースコード

入力された頂点 node に対して \mathrm{dp}_{\mathrm{node}} を求める関数 dfs() を、以下のように実装しました。

int l[int(3e5 + 5)], r[int(3e5 + 5)];

int dp[int(3e5 + 5)];
bool visited[int(3e5 + 5)]; // visited は dp[node] を既に計算していたら true にする

int dfs(int node) {
  if (visited[node]) {
    return dp[node]; // もし既に dp[node] の値を計算していたら、そのまま dp[node] を返す
  }

  visited[node] = true; // dp[node] の値を計算したことにする

  if (l[node] == 0 && r[node] == 0) {
    return dp[node] = 0; // もし子が存在しないならば、 dp[node] を 0 として更新して値を返す
  }

  int ans = 1e9; // 答えとなる値。一旦 10^9 にしておく
  if (l[node] != 0) { // 左側の子が存在する場合
    if (s[node - 1] == 'L') { // s は文字列で 0-index なので注意する
      ans = min(ans, dfs(l[node])); // もし s[i] = L であれば、ans と dp[l[node]] の値のうち最小値が ans となる
    } else {
      ans = min(ans, dfs(l[node]) + 1); // そうでないならば、 ans と dp[l[node]] + 1 の値のうち最小値が ans となる
    }
  }
  if (r[node] != 0) { // 右側の子が存在する場合
    if (s[node - 1] == 'R') {
      ans = min(ans, dfs(r[node])); // もし s[i] = R であれば、ans と dp[r[node]] の値のうち最小値が ans となる
    } else {
      ans = min(ans, dfs(r[node]) + 1); // そうでないならば、 ans と dp[r[node]] + 1 の値のうち最小値が ans となる
    }
  }

  return dp[node] = ans; // dp[node] を ans として更新して値を返す
}

感想

前回 に続いて、また再帰の考え方を使いました。（といっても、再帰といえるほど繰り返していないと思います。）

再帰については「終着点をどのようにするか」と「答えを計算したかどうかをどう管理するか」の2つのポイントが個人的には大切だと思っていて、ここをしっかり定義できていれば、実装にも迷いが出なくなると思います。