AtCoder Beginner Contest 177 参加記

問題文は https://atcoder.jp/contests/abc177 こちら。
Rustで解いています。https://github.com/tanakh/competitive-rs このライブラリを使っています。
入出力はライブラリで自動化されています。

A - Don't be late

距離DをスピードSで移動するのにかかる時間がT以下かどうかを判別する問題。

$D/S≦T$ でいいんだけど、浮動小数点の誤差を考えたくないので両辺にSを掛けて $D≦TS$ にするのが楽。
掛け算をすると計算結果がでかくなるので、オーバーフローしないかはちゃんとチェックしておきたい。
今回の場合は $1≦D,T,S≦10,000$ なので、$TS$ は32ビット整数でも収まる範囲なので問題ない。
下のコードでは、雑に i64 を使ってます。

use competitive::prelude::*;

# [argio(output = AtCoder)]
fn main(d: i64, t: i64, s: i64) -> bool {
    d <= t * s
}

B - Substring

2つの文字列$S$と$T$ $(1≦|S|,|T|≦1000, |T|≦|S|)$が与えられて、$T$が$S$の部分文字列になるように$S$を書き換えるとき、少なくとも何文字書き換える必要があるか？という問題。

$T$を$S$のどこに一致させるかを決めれば、その時に必要になる書き換え回数は単に一致しない文字の数である。
すべての位置を試して、最小値を取るというナイーブなアルゴリズムだと計算量は $O(|S||T|)$ だけど、
$|S|$も$|T|$も小さいのでこれで十分間に合う。

下のコードでは、すべての開始位置iに対して、s[i..]とtをzipして、!= であるようなところの数を数えるみたいな処理を書いている。

use competitive::prelude::*;

# [argio(output = AtCoder)]
fn main(s: Chars, t: Chars) -> usize {
    (0..s.len() - t.len() + 1)
        .map(|i| s[i..].iter().zip(&t).filter(|(c, d)| c != d).count())
        .min()
        .unwrap()
}

C - Sum of product of pairs

$A_1, ... A_N (0≦A_i≦10^9, N≦210^5)$ に対して、$\sum_{1≦i<j≦N}{A_iA_j}$ を求める問題。結果は $1,000,000,009$ で割ったあまりを返す。

ナイーブに$O(N^2)$回の積和を行うとさすがに2秒では間に合わない。SIMDでもやろうとしてみたけどどうにもうまくいかなかった。x86では64ビット整数の乗算がSIMDでできないのが痛い気がする。f64で53ビット整数にしてもこの問題の場合はあんまり意味がないし。SIMDで2乗のコードが通せた人がいたら教えてほしい。

$i<j$の条件を無視して考えると、$A_i$を左側に含む項は $A_iA_1+...+A_iA_N = A_i*(A_1+...+A_N)= A_i*\sum{A}$ になるので、これらを$A$の全要素について足し合わせたものは $A_1*\sum{A}+...+A_N*\sum{A}=(A_1+...+A_N)*\sum{A} = \sum{A} * \sum{A}$ になる。

これは求めたい値に対して $i=j$ と $i>j$ についての項を足し過ぎているので、これらを引いてやる。前者は $\sum_i^N{A_i*A_i}$ だし、後者は乗算の交換法則から $i<j$ の場合と同じになるはずなので、結局$i=j$の部分を引いてから2で割ればいい。

下のコードでは、合計の値が $N*max(A)^2$ 以下になるはずで、これはi128の範囲に収まるというのを利用して、i128で計算してから最後に剰余を取っている。せっかくi128があるのにAtCoderではあんまり使える機会がないから使ってあげてみた。

use competitive::prelude::*;

# [argio(output = AtCoder)]
fn main(n: usize, a: [i128; n]) -> i128 {
    (a.iter().sum::<i128>().pow(2) - a.iter().map(|a| a * a).sum::<i128>()) / 2 % 1_000_000_007
}

D - Friends

N(≦2*10^5)人の人がいて、友達の友達は友達みたいな友達関係が与えられるから、このN人をそれぞれ友達同士を含まないグループに分割したいんだけど、最小でいくつのグループが必要になるかという問題。

Union Find で友人グループを求めると、各友人グループに属する人は同じグループには属せないので、一番でかい友人グループの人数が答えになる。

use competitive::prelude::*;
use competitive::union_find::UnionFind;

# [argio(output = AtCoder)]
fn main(n: usize, m: usize, ab: [(Usize1, Usize1); m]) -> usize {
    let mut uf = UnionFind::new(n);
    for (a, b) in ab {
        uf.union(a, b);
    }

    let mut mm = BTreeMap::<usize, usize>::new();
    for i in 0..n {
        *mm.entry(uf.find(i)).or_default() += 1;
    }
    *mm.values().max().unwrap()
}

E - Coprime

N(≦10^6)個の整数$A_i(≦10^6)$が与えられる時、それらの任意の2つの数の最大公約数が1の時は "pairwise coprime"、すべての数の最大公約数が1の時は "setwise coprime"、どちらでもないときは "not coprime" を出力するという問題。

setwise coprimeの判定はgcdでfoldすればいいだけなので、実質 "pairwise coprime" を判定する問題。

これまたナイーブにやるとNがでかいので制限時間的に難しい。ある数 $A_i$ に対して、それがほかの数すべてと互いに素であるというのがどういうことかって考えてみると、つまり $A_i$のすべての素因数がほかの数に素因数として含まれないということになる。すべての素数に対してあらかじめそれを素因数として含む数のリストを全部計算しておけば、これは$A_i$の素因数の個数程度の計算量でチェックできる。ある数の素因数の数は明らかに$O(logN)$を超えないので、これは十分高速なはず。

use competitive::prelude::*;

fn factor(n: usize, ps: &[usize]) -> Vec<(usize, usize)> {
    let mut n = n;
    let mut ret = vec![];
    for &p in ps {
        let mut cnt = 0;
        while n % p == 0 {
            n /= p;
            cnt += 1;
        }
        if cnt > 0 {
            ret.push((p, cnt));
        }
    }
    if n > 1 {
        ret.push((n, 1));
    }
    ret
}

# [argio(output = AtCoder)]
fn main(n: usize, a: [usize; n]) -> &str {
    let mut prime_pos = BTreeMap::<usize, Vec<usize>>::new();
    let ps = competitive::prime::primes(1001);

    let factors: Vec<Vec<(usize, usize)>> = a.iter().map(|a| factor(*a, &ps)).collect();

    for (i, fs) in factors.iter().enumerate() {
        for &(p, _) in fs.iter() {
            prime_pos.entry(p).or_default().push(i);
        }
    }

    let mut pairwise = true;

    for (i, fs) in factors.iter().enumerate() {
        let mut not_coprime = false;
        for &(p, _) in fs.iter() {
            for &j in prime_pos.get(&p).unwrap().iter() {
                if i == j {
                    continue;
                }
                not_coprime = true;
                break;
            }

            if not_coprime {
                break;
            }
        }

        if not_coprime {
            pairwise = false;
            break;
        }
    }

    if pairwise {
        return "pairwise coprime";
    }

    if a.iter().fold(0, |i, &j| gcd(i, j)) == 1 {
        return "setwise coprime";
    }

    "not coprime"
}

このコードは汎用素因数分解コードをそのまま使って超ナイーブに素因数分解しているが、これは1.6秒ぐらいかかってて結構TLEぎりぎりであやうい。素因数分解したい値が小さい（≦1e6）ので、エラトステネスのふるいの要領で計算しながら素因数を記録していけば速くなる。

use competitive::prelude::*;

# [argio(output = AtCoder)]
fn main(n: usize, a: [usize; n]) -> &str {
    const MAX_A: usize = 1_000_001;

    // エラトステネスの篩で2..MAX_Aまでの素因数分解を一気に全部作る
    let mut factors = vec![vec![]; MAX_A];
    for i in 2..factors.len() {
        if factors[i].is_empty() {
            for j in (i..factors.len()).step_by(i) {
                factors[j].push(i);
            }
        }
    }

    // 素因数からそれを含む数へのマップを作る
    let mut prime_pos = vec![vec![]; MAX_A];
    for i in 0..n {
        for &p in factors[a[i]].iter() {
            prime_pos[p].push(i);
        }
    }

    // 0..nのすべてのa[i]に対して、それの素因数すべてについて、それが含まれる数がiのみである
    if (0..n).all(|i| factors[a[i]].iter().all(|&p| prime_pos[p] == &[i])) {
        "pairwise coprime"
    } else if a.iter().fold(0, |i, &j| gcd(i, j)) == 1 {
        "setwise coprime"
    } else {
        "not coprime"
    }
}

これで0.7秒弱になった。まだちょっと時間がかかり気味に見えるけど、リファクタリングしたらなんだかコードがすごくすっきりしたので満足。

F - I hate Shortest Path Problem

$W×H(W,H≦2*10^5)$ の2次元の盤面上に障害物があって右と下にしか行けないときに、各行のいずれかの位置に到着するための最小の移動回数を求める問題。

盤面はがでかいので愚直には計算できない。障害物が上下方向を遮るようにしか置かれていないので、右方向への移動は常に可能。ということは、右はいつでも行けるので、下に行けるときは常に下に行けばいいということになる。一番上の行のすべての場所からスタートすると、右に移動が入ると別の場所からスタートしたものと合流することがある。そして、合流したらその後はずっと同じように移動するはず。

なので、各行において、どのカラムに、最短横移動回数いくつでたどり着けるのかというのを保持しながら、1行ずつシミュレートしていけばよくて、シミュレートは一見重そうに見えるけど、下に障害物がないものについては、そのまま下に移動させればよくて、この場合は横移動回数が増えないので何もしなくていい。下に障害物があるときは、その上に載ってる奴を全部その障害物の右端まで移動させて距離の最小値を取らなければならない。これは乗っている個数に比例した時間がかかるが、移動させた後はすべて合流して1つになる。つまり移動はO(1)だし、横移動は横に移動させたものの数に比例する時間がかかるが、移動させただけ管理する位置の数が減っていく。要するに、ワーストケースでO(H+W)程度の時間しかかからないはずなのでこれは十分早く計算できる。

各行における最小の移動回数がどれかを探さなければならないが、これを毎回全部調べていたらこれは間に合わない。なので、multisetなどを使ってどの移動回数が何か所あるかを管理していく。そうすればO(logW)ぐらいの時間で各行に対する答えが求まる。

use competitive::prelude::*;

# [argio(output = AtCoder)]
fn main(h: usize, w: usize, ab: [(Usize1, Usize1); h]) -> Vertical<Option<usize>> {
    // ssは現在の行における到達可能位置とそこへの最短移動回数の集合
    // msは最短移動回数のmultiset
    let mut ss = BTreeSet::new();
    let mut ms = MultiSet::new();

    for i in 0..w {
        ss.insert((i, 0));
        ms.insert(0);
    }

    let mut ret = vec![];

    for (i, (l, r)) in ab.into_iter().enumerate() {
        let r = r + 1;
        let mut rem = vec![];
        let mut minv = usize::max_value();

        // 行iの障害物の上に載っているものを列挙する
        for &(c, d) in ss.range((l, 0)..(r, 0)) {
            if r != w {
                let cv = d + r - c;
                minv = min(minv, cv);
            }
            rem.push((c, d));
        }

        // 障害物の上に載っていたものを消して、
        for rem in rem {
            ms.remove(&rem.1);
            ss.remove(&rem);
        }

        // 障害物の一個右の位置に移動させる（右端がWなら下の段に移動できないので消滅）
        if minv < usize::max_value() {
            ss.insert((r, minv));
            ms.insert(minv);
        }

        ret.push(ms.min().map(|r| (i + 1 + r)));
    }

    Vertical(ret)
}