AtCoder精選良問「B - Values」Diff 732

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

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

問題リンク

B - Values

問題概要

N頂点M辺の単純無向グラフで、各辺は二つの頂点を結んでいる。
はじめ、頂点 i には値 ai​ が書かれている。
そのような設定で以下の操作を行う。

操作：一つの辺を選び、その辺が結ぶ頂点の一つの値を1増やし、もう一つの値を1減らす、またはその逆の操作を行う。

操作後に各頂点の値をそれぞれ b1​,b2,⋯⋯,bN​にすることができるか、出力する問題。

考え方

まず、配列を２つ考えるのは見通しが悪いので、差分を取って操作しなければいけない数字の配列を作る。

1 2 3
2 2 2

入力例１の場合はこんな数列が与えられるので、B-Aで差分を取ると、以下のようになる。

1 0 -1 

とりあえず、この差分を適切に操作して全て０にすることを聞いている問題なんだな、ということを理解する。

次に操作に目を向ける。
操作は二つの頂点を選んで片方を-1,片方を+1するような操作だった。
重要な考察として、操作をした数列の総和自体は変化しない。

ということは、数列の総和が０でない時点で、絶対に全て０にすることはできない。

逆に考えると「総和が０だったら絶対にできるんじゃないか」と考える。

例えばこんな数列があったとする。

5 -7 2 -3 3

総和は０。
左から貪欲に操作すると絶対に０にできる。

// 1回目
0 -2 2 -3 3
// 2回目
0 0 0 -3 3
// 3回目
0 0 0 0 0

グラフ自体は単純グラフであるが、いくつかのグループに分かれることがあるようだ。入力例３のグラフをプロットするとこんな感じ。

ということは各グループごとの差分配列の総和が０になるか考えればよい。
これはUnionFindなどのデータ構造を使うと比較的簡単に実装できる。

実装

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

export class UnionFind {
    private parents: number[];

    constructor(private n: number) {
        this.parents = Array.from({ length: n }, () => -1);
    }

    public find(x: number): number {
        if (this.parents[x] < 0) {
            return x;
        } else {
            this.parents[x] = this.find(this.parents[x]);
            return this.parents[x];
        }
    }

    public union(x: number, y: number): void {
        x = this.find(x);
        y = this.find(y);

        if (x === y) {
            return;
        }

        if (this.parents[x] > this.parents[y]) {
            [x, y] = [y, x];
        }

        this.parents[x] += this.parents[y];
        this.parents[y] = x;
    }

    public size(x: number): number {
        return -this.parents[this.find(x)];
    }

    public same(x: number, y: number): boolean {
        return this.find(x) === this.find(y);
    }

    public members(x: number): number[] {
        const root = this.find(x);
        return Array.from({ length: this.n }, (_, i) => i).filter(i => this.find(i) === root);
    }

    public roots(): number[] {
        return Array.from({ length: this.n }, (_, i) => i).filter(i => this.parents[i] < 0);
    }

    public groupCount(): number {
        return this.roots().length;
    }

    public allGroupMembers(): Map<number, number[]> {
        const groupMembers = new Map<number, number[]>();
        for (let member = 0; member < this.n; member++) {
            const root = this.find(member);
            if (!groupMembers.has(root)) {
                groupMembers.set(root, []);
            }
            groupMembers.get(root)!.push(member);
        }
        return groupMembers;
    }
}

function main() {
    const [n, m] = nextNums(2)
    const arrA = nextNums(n)
    const arrB = nextNums(n)
    // 差を取る
    const diff: number[] = []
    for (let i = 0; i < n; i++) {
        diff.push(arrB[i] - arrA[i])
    }
    // UnionFindを作成
    const uf = new UnionFind(n)
    // 辺ごとに頂点をグルーピングする
    for (let i = 0; i < m; i++) {
        let [a, b] = nextNums(2)
        // 0インデックスに直す
        a--
        b--
        // unionする
        uf.union(a, b)
    }
    let ans = 'Yes'
    // 各メンバーを繰り返す
    for (const elm of uf.allGroupMembers()) {
        const members = elm[1]
        let arrSum = 0
        for (const vertex of members) {
            arrSum += diff[vertex]
        }
        // 0じゃなかったらできない
        if (arrSum != 0) {
            ans = 'No'
        }
    }
    log(ans)
}

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();
}