AtCoder精選良問「D - Coloring Edges on Tree」Diff 1192

この記事はAtCoderの精選良問集マガジンの中の記事です。
私が良問だと思った問題を解説付きでまとめています。

AtCoder 精選良問集 Diff 茶後半～水前半｜Qlitre｜note

問題リンク

D - Coloring Edges on Tree

問題概要

N個の頂点からなる木Gが与えられる。
各頂点には1からNまでの番号がついており、i番目の辺は頂点a_iと頂点b_iを結んでいます。Gの辺を何色かで塗り分けることを考える。
各頂点について、その頂点を端点に持つ辺の色がすべて異なるようにしたいです。
この条件を満たす塗り分けの中で、使用する色の数が最小であるようなものを1つ構築する。

考え方

難しい。問題から一見して幅優先探索や深さ優先探索を使うように思えるが、辺の管理をしたことは今までなかった。

まずは、入力例２のデータを絵に書いてみる。

8
1 2
2 3
2 4
2 5
4 7
5 6
6 8

制約で1≤ai​<bi​≤Nとなっているので、木の根は必ず１になる。枝わかれしていくにつれて、段々と数が大きくなっていく、という構造。

シンプルにかぶらないように色を塗っていったら正解っぽい気がする。
色は１から始まるので、うえから順番に１から塗っていって、かぶってたら、数字を増やすという戦略で塗ってみる。

貪欲に１から始めて塗れたら塗る、塗れなかったら数字をカウントアップする、という方法で塗ると最小の色で塗り分けができそう。

あとはこれを実装する方法を考える。

通常の探索だと、点を辿っていく。
当たり前であるが、点を辿るときにはかならず辺を通ることになる。
なので、このときに答えを記録していくような感じで実装をする。グラフを構築する際に辺の番号も保持しておくとやりやすい。
また、探索中には「今見ている点に来た際にどの色を使ったか」も管理をする。
そうして塗り分けをする際に数字が被らないようにする。

実装

import { createInterface } from "readline";
import * as fs from "fs";
import { log } from "console";
import { exit } from "process";

function main() {
    const n = nextNum()
    // 木構造。[[頂点,辺],[頂点,辺]..]という構造にする。
    const graph: Record<number, number[][]> = {}
    // 初期化
    for (let i = 1; i < n + 1; i++) {
        graph[i] = []
    }
    // 木を構築する
    for (let i = 0; i < n - 1; i++) {
        const [a, b] = nextNums(2)
        graph[a].push([b, i])
    }
    const stack: number[][] = []
    // [頂点,辿ってきた色]
    stack.push([1, 0])
    const ans: number[] = new Array(n - 1).fill(n)
    while (stack.length > 0) {
        const [vertex, parentColor] = stack.pop()!
        // 答えとなる色。貪欲に１から始める
        let color = 1
        const chirdlen = graph[vertex]
        // 子頂点に遷移する
        for (const elm of chirdlen) {
            const [child, edgeNum] = elm
            // 親のカラーと一緒だったらカウントアップ
            if (color == parentColor) color++
            // 色を塗る
            ans[edgeNum] = color
            // stackに加える
            stack.push([child, color])
            // 今の色は使えないのでカウントアップする
            color++
        }
    }
    // 最大の色を調べる
    let k = 0
    for (const color of ans) {
        k = Math.max(k, color)
    }
    // 答えの出力
    log(k)
    for (const color of ans) {
        log(color)
    }
}

let inputs = "";
let inputArray: string[];
let currentIndex = 0;

function next() {
    return inputArray[currentIndex++];
}
function nextNum() {
    return +next();
}
function nextBigInt() {
    return BigInt(next());
}
function nexts(length: number) {
    const arr = [];
    for (let i = 0; i < length; ++i) arr[i] = next();
    return arr;
}
function nextNums(length: number) {
    const arr = [];
    for (let i = 0; i < length; ++i) arr[i] = nextNum();
    return arr;
}
function nextBigInts(length: number) {
    const arr = [];
    for (let i = 0; i < length; ++i) arr[i] = nextBigInt();
    return arr;
}


// デバッグ環境がWindowsであれば条件分岐する
if (process.env.OS == "Windows_NT") {
    const stream = createInterface({
        input: process.stdin,
        output: process.stdout,
    });
    stream.on("line", (line) => {
        inputs += line;
        inputs += "\n";
    });
    stream.on("close", () => {
        inputArray = inputs.split(/\s/);
        main();
    });
} else {
    inputs = fs.readFileSync("/dev/stdin", "utf8");
    inputArray = inputs.split(/\s/);
    main();
}