2019-10-06

AGC039 C - Division by Two with Something

問題

解法

各操作を観察すると、どちらも「最も下のビットをpopして反転して最も上にpushする」と言い換えられる。これをN回繰り返すと元の数のビットを反転させたものが得られるので、2N回繰り返すと必ず元の数が得られる。よって繰り返して元の数が得られる繰り返し回数kは2Nの約数であると考えられる。

繰り返し回数をkとし、繰り返されている長さkのビット列を[..., c, b, a] とすると、k回操作した後のビット列は[..., !c, !b, !a, ..., c, b, a, ..., c, b, a, ..., ..., c, b, a] という感じになっている。ここでNがkの倍数だとするとN回繰り返した後に元の数が得られるはずだが、N回操作後のビット列は[..., c!, b!, !a]となっていて元に戻ってはいないため、矛盾する。よってkはNの約数ではない。kは2Nの約数ではあるため、kは偶数であり、2N/kは奇数である。

2N/kが奇数であることから長さkのビット列は[..., !c, !b, !a, ..., c, b, a]となっているはずである。よってこの長さkのビット列は長さk/2のビット列と、それを反転したものを並べたものだと分かる。

各kについて動的計画法で数え上げを計算し、重複したものを包除原理で取り除いてやることで、答えが求まる。

コード

use mod_int::ModInt;
use std::collections::BTreeMap;

const MOD: usize = 998244353;
fn main() {
    let s = std::io::stdin();
    let mut sc = Scanner { stdin: s.lock() };
    let n: usize = sc.read();
    let x = sc
        .chars()
        .into_iter()
        .map(|c| if c == '0' { 0 } else { 1 })
        .collect::<Vec<_>>();

    let m = 2 * n;
    let mut ks = vec![];
    for k in 2..(m + 1) {
        if m % k == 0 && k % 2 == 0 && (m / k) % 2 == 1 {
            ks.push(k);
        }
    }

    let mut ans = vec![ModInt(0); m + 1];
    for &k in ks.iter() {
        let mut dp = vec![vec![ModInt(0); 2]; 2];
        dp[0][0] = ModInt(1);
        for i in 0..(k / 2) {
            let mut next_dp = vec![vec![ModInt(0); 2]; 2];
            for prev in 0..2 {
                for next in 0..2 {
                    for smaller in 0..2 {
                        if smaller == 0 && next > x[i] {
                            continue;
                        }
                        let next_smaller = if next < x[i] { 1 } else { smaller };
                        next_dp[next][next_smaller] += dp[prev][smaller];
                    }
                }
            }
            dp = next_dp;
        }

        let same = dp[0][0] + dp[1][0];
        let smaller = dp[0][1] + dp[1][1];
        let mut ok = true;
        for i in 0..n {
            let s = i % k;
            let t = i % (k / 2);
            let y = if s == t {
                x[i % (k / 2)]
            } else {
                x[i % (k / 2)] ^ 1
            };
            if x[i] > y {
                break;
            }
            if x[i] < y {
                ok = false;
                break;
            }
        }
        ans[k] = if ok { same + smaller } else { smaller };
    }

    let mut result = ModInt(0);
    for k in 1..(m + 1) {
        let mut cur = k * 2;
        while cur <= m {
            if ans[cur].0 > 0 {
                ans[cur] -= ans[k];
            }
            cur += k;
        }
        result += ans[k] * k;
    }
    println!("{}", result.0);
}

pub mod mod_int {
    use super::MOD;
    use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Sub, SubAssign};

    type Num = usize;

    #[derive(Clone, Copy, Debug)]
    pub struct ModInt<T: Copy + Clone>(pub T);

    impl Add<ModInt<Num>> for ModInt<Num> {
        type Output = ModInt<Num>;
        fn add(self, rhs: ModInt<Num>) -> ModInt<Num> {
            self + rhs.0
        }
    }

    impl Add<Num> for ModInt<Num> {
        type Output = ModInt<Num>;
        fn add(self, rhs: Num) -> ModInt<Num> {
            let mut t = rhs + self.0;
            if t >= MOD {
                t = t - MOD;
            }
            ModInt(t)
        }
    }

    impl Sub<Num> for ModInt<Num> {
        type Output = ModInt<Num>;
        fn sub(self, rhs: Num) -> ModInt<Num> {
            let rhs = if rhs >= MOD { rhs % MOD } else { rhs };
            let value = if self.0 < rhs { self.0 + MOD } else { self.0 };
            ModInt(value - rhs)
        }
    }

    impl Sub<ModInt<Num>> for ModInt<Num> {
        type Output = ModInt<Num>;
        fn sub(self, rhs: ModInt<Num>) -> ModInt<Num> {
            self - rhs.0
        }
    }

    impl AddAssign<Num> for ModInt<Num> {
        fn add_assign(&mut self, other: Num) {
            *self = *self + other;
        }
    }
    impl AddAssign<ModInt<Num>> for ModInt<Num> {
        fn add_assign(&mut self, other: ModInt<Num>) {
            *self = *self + other;
        }
    }

    impl SubAssign<Num> for ModInt<Num> {
        fn sub_assign(&mut self, other: Num) {
            *self = *self - other;
        }
    }

    impl SubAssign<ModInt<Num>> for ModInt<Num> {
        fn sub_assign(&mut self, other: ModInt<Num>) {
            *self = *self - other;
        }
    }

    impl Div<Num> for ModInt<Num> {
        type Output = ModInt<Num>;
        fn div(self, rhs: Num) -> ModInt<Num> {
            self * ModInt(rhs).pow(MOD - 2)
        }
    }

    impl Div<ModInt<Num>> for ModInt<Num> {
        type Output = ModInt<Num>;
        fn div(self, rhs: ModInt<Num>) -> ModInt<Num> {
            self / rhs.0
        }
    }

    impl DivAssign<Num> for ModInt<Num> {
        fn div_assign(&mut self, rhs: Num) {
            *self = *self / rhs
        }
    }
    impl DivAssign<ModInt<Num>> for ModInt<Num> {
        fn div_assign(&mut self, rhs: ModInt<Num>) {
            *self = *self / rhs
        }
    }

    impl Mul<ModInt<Num>> for ModInt<Num> {
        type Output = ModInt<Num>;

        fn mul(self, rhs: ModInt<Num>) -> ModInt<Num> {
            self * rhs.0
        }
    }
    impl Mul<Num> for ModInt<Num> {
        type Output = ModInt<Num>;

        fn mul(self, rhs: Num) -> ModInt<Num> {
            let t = (self.0 * rhs) % MOD;
            ModInt(t)
        }
    }

    impl MulAssign<Num> for ModInt<Num> {
        fn mul_assign(&mut self, rhs: Num) {
            *self = *self * rhs;
        }
    }

    impl MulAssign<ModInt<Num>> for ModInt<Num> {
        fn mul_assign(&mut self, rhs: ModInt<Num>) {
            *self = *self * rhs;
        }
    }

    impl ModInt<Num> {
        pub fn pow(self, e: usize) -> ModInt<Num> {
            let mut result = ModInt(1);
            let mut cur = self;
            let mut e = e;
            while e > 0 {
                if e & 1 == 1 {
                    result *= cur;
                }
                e >>= 1;
                cur *= cur;
            }
            result
        }
    }
}

pub struct Scanner<R> {
    stdin: R,
}

impl<R: std::io::Read> Scanner<R> {
    pub fn read<T: std::str::FromStr>(&mut self) -> T {
        use std::io::Read;
        let buf = self
            .stdin
            .by_ref()
            .bytes()
            .map(|b| b.unwrap())
            .skip_while(|&b| b == b' ' || b == b'\n')
            .take_while(|&b| b != b' ' && b != b'\n')
            .collect::<Vec<_>>();
        unsafe { std::str::from_utf8_unchecked(&buf) }
            .parse()
            .ok()
            .expect("Parse error.")
    }
    pub fn vec<T: std::str::FromStr>(&mut self, n: usize) -> Vec<T> {
        (0..n).map(|_| self.read()).collect()
    }
    pub fn chars(&mut self) -> Vec<char> {
        self.read::<String>().chars().collect()
    }
}

2019-09-26

AGC038 C - LCMs

AtCoder Rust 競技プログラミング数学

問題

atcoder.jp

解法

mobile.twitter.com

自分の学習のために式を書き直してみたが全く同じになってしまった……

$\begin{eqnarray} \sum_{i,j} LCM(A_i, A_j) & = & \sum_{i,j} \frac{A_i,A_j}{GCD(A_i, A_j)} \\ & = & \sum_{d} \left(\frac{1}{d} \sum _{i,j, GCD(A_i,A_j)=d}A_i A_j \right)\\ & = & \sum_{d} \left( f(d) \sum _{i,j,d | GCD(A_i,A_j)}A_i A_j \right) \\ & = & \sum_{d} \left( f(d)\sum_{d | x} \left(\sum _{i,j, GCD(A_i,A_j)=x}A_i A_j \right) \right) \\ \end{eqnarray}$

ここで
$\begin{eqnarray} F(x) &=& \sum _{i,j, GCD(A_i,A_j)=x}A_i A_j \end{eqnarray}$
とすると
$\begin{eqnarray} \sum_{d} (\frac{1}{d} F(d) ) & =& \sum_{d}(f(d)\sum_{d|x} F(x)) \\ & =& \sum_{x}(F(x)\sum_{d|x} f(d)) \\ & =& \sum_{d}(F(d)\sum_{x|d} f(x)) \end{eqnarray}$
なので
$\begin{eqnarray} \sum_{x|d} f(x) &=& \frac{1}{d} \end{eqnarray}$

コード

use mod_int::ModInt;

const MOD: usize = 998244353;

fn main() {
    let s = std::io::stdin();
    let mut sc = Scanner { stdin: s.lock() };
    let n: usize = sc.read();
    let a: Vec<usize> = sc.vec(n);
    let max_a = *a.iter().max().unwrap();
    let mut f = vec![ModInt(0); max_a + 1];
    for i in 1..(max_a + 1) {
        f[i] = ModInt(1) / i;
    }
    for x in 1..(max_a + 1) {
        let mut d = x * 2;
        while d <= max_a {
            let fx = f[x];
            f[d] -= fx;
            d += x;
        }
    }

    let mut count = vec![0; max_a + 1];
    for &a in a.iter() {
        count[a] += 1;
    }

    let mut ans = ModInt(0);
    for d in 1..(max_a + 1) {
        let mut x = d;
        let mut sum = ModInt(0);
        let mut sum2 = ModInt(0);
        while x <= max_a {
            // x%d == 0
            // d|x
            sum += ModInt(x) * count[x];
            sum2 += ModInt(x) * x * count[x];
            x += d;
        }
        ans += (sum * sum - sum2) / 2 * f[d];
    }
    println!("{}", ans.0);
}

pub mod mod_int {
    use super::MOD;
    use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Sub, SubAssign};

    type Num = usize;

    #[derive(Clone, Copy)]
    pub struct ModInt<T: Copy + Clone>(pub T);

    impl Add<ModInt<Num>> for ModInt<Num> {
        type Output = ModInt<Num>;
        fn add(self, rhs: ModInt<Num>) -> ModInt<Num> {
            self + rhs.0
        }
    }

    impl Add<Num> for ModInt<Num> {
        type Output = ModInt<Num>;
        fn add(self, rhs: Num) -> ModInt<Num> {
            let mut t = rhs + self.0;
            if t >= MOD {
                t = t - MOD;
            }
            ModInt(t)
        }
    }

    impl Sub<Num> for ModInt<Num> {
        type Output = ModInt<Num>;
        fn sub(self, rhs: Num) -> ModInt<Num> {
            let rhs = if rhs >= MOD { rhs % MOD } else { rhs };
            let value = if self.0 < rhs { self.0 + MOD } else { self.0 };
            ModInt(value - rhs)
        }
    }

    impl Sub<ModInt<Num>> for ModInt<Num> {
        type Output = ModInt<Num>;
        fn sub(self, rhs: ModInt<Num>) -> ModInt<Num> {
            self - rhs.0
        }
    }

    impl AddAssign<Num> for ModInt<Num> {
        fn add_assign(&mut self, other: Num) {
            *self = *self + other;
        }
    }
    impl AddAssign<ModInt<Num>> for ModInt<Num> {
        fn add_assign(&mut self, other: ModInt<Num>) {
            *self = *self + other;
        }
    }

    impl SubAssign<Num> for ModInt<Num> {
        fn sub_assign(&mut self, other: Num) {
            *self = *self - other;
        }
    }

    impl SubAssign<ModInt<Num>> for ModInt<Num> {
        fn sub_assign(&mut self, other: ModInt<Num>) {
            *self = *self - other;
        }
    }

    impl Div<Num> for ModInt<Num> {
        type Output = ModInt<Num>;
        fn div(self, rhs: Num) -> ModInt<Num> {
            self * ModInt(rhs).pow(MOD - 2)
        }
    }

    impl Div<ModInt<Num>> for ModInt<Num> {
        type Output = ModInt<Num>;
        fn div(self, rhs: ModInt<Num>) -> ModInt<Num> {
            self / rhs.0
        }
    }

    impl DivAssign<Num> for ModInt<Num> {
        fn div_assign(&mut self, rhs: Num) {
            *self = *self / rhs
        }
    }
    impl DivAssign<ModInt<Num>> for ModInt<Num> {
        fn div_assign(&mut self, rhs: ModInt<Num>) {
            *self = *self / rhs
        }
    }

    impl Mul<ModInt<Num>> for ModInt<Num> {
        type Output = ModInt<Num>;

        fn mul(self, rhs: ModInt<Num>) -> ModInt<Num> {
            self * rhs.0
        }
    }
    impl Mul<Num> for ModInt<Num> {
        type Output = ModInt<Num>;

        fn mul(self, rhs: Num) -> ModInt<Num> {
            let t = (self.0 * rhs) % MOD;
            ModInt(t)
        }
    }

    impl MulAssign<Num> for ModInt<Num> {
        fn mul_assign(&mut self, rhs: Num) {
            *self = *self * rhs;
        }
    }

    impl MulAssign<ModInt<Num>> for ModInt<Num> {
        fn mul_assign(&mut self, rhs: ModInt<Num>) {
            *self = *self * rhs;
        }
    }

    impl ModInt<Num> {
        pub fn pow(self, e: usize) -> ModInt<Num> {
            let mut result = ModInt(1);
            let mut cur = self;
            let mut e = e;
            while e > 0 {
                if e & 1 == 1 {
                    result *= cur;
                }
                e >>= 1;
                cur *= cur;
            }
            result
        }
    }
}

pub struct Scanner<R> {
    stdin: R,
}

impl<R: std::io::Read> Scanner<R> {
    pub fn read<T: std::str::FromStr>(&mut self) -> T {
        use std::io::Read;
        let buf = self
            .stdin
            .by_ref()
            .bytes()
            .map(|b| b.unwrap())
            .skip_while(|&b| b == b' ' || b == b'\n')
            .take_while(|&b| b != b' ' && b != b'\n')
            .collect::<Vec<_>>();
        unsafe { std::str::from_utf8_unchecked(&buf) }
            .parse()
            .ok()
            .expect("Parse error.")
    }
    pub fn vec<T: std::str::FromStr>(&mut self, n: usize) -> Vec<T> {
        (0..n).map(|_| self.read()).collect()
    }
    pub fn chars(&mut self) -> Vec<char> {
        self.read::<String>().chars().collect()
    }
}

2019-09-24

Project Euler 679 Freefarea

Project Euler 競技プログラミング数学 Rust

問題

projecteuler.net

長さ n で A, E, F, R からなる文字列のうち、4つのキーワード FREE, FARE, AREA, REEF を各1回ずつ含むものはいくつあるか。n=30について求めよ。

解法

全ての文字列を陽に持つ必要はなく、(各キーワードを持っているかどうか, 接尾辞4文字, 個数) を持ちながら伸ばしていけば良い。

コード

use std::collections::BTreeMap;

const A: u64 = 0;
const E: u64 = 1;
const F: u64 = 2;
const R: u64 = 3;

const LETTERS: [u64; 4] = [A, E, F, R];

const FREE: u64 = (F << 6) | (R << 4) | (E << 2) | E;
const FARE: u64 = (F << 6) | (A << 4) | (R << 2) | E;
const AREA: u64 = (A << 6) | (R << 4) | (E << 2) | A;
const REEF: u64 = (R << 6) | (E << 4) | (E << 2) | F;

const KEYWORDS: [u64; 4] = [FREE, FARE, AREA, REEF];

fn main() {
    let mut words = BTreeMap::new();
    words.insert((0, 0), 1);
    for length in 1..31 {
        let mut next = BTreeMap::new();
        for ((mask, suffix), count) in words.into_iter() {
            for &letter in LETTERS.iter() {
                let next_word = (suffix << 2) | letter;
                let next_suffix = next_word & ((1 << 8) - 1);

                match KEYWORDS
                    .iter()
                    .enumerate()
                    .find(|&(_, &keyword)| keyword == next_suffix)
                    .map(|(i, _)| i)
                {
                    Some(new_word_id) if length >= 4 => {
                        if mask & (1 << new_word_id) == 0 {
                            let next_mask = mask | (1 << new_word_id);
                            *next.entry((next_mask, next_suffix)).or_insert(0) += count;
                        }
                    }
                    _ => {
                        *next.entry((mask, next_suffix)).or_insert(0) += count;
                    }
                }
            }
        }

        let ans = next
            .iter()
            .filter(|&(&(mask, _), _)| mask == (1 << 4) - 1)
            .map(|(_, &count)| count)
            .sum::<usize>();
        words = next;
        eprintln!("length={} ans={}", length, ans);
    }
    //    naive(11);
}

fn decode_suffix(suffix: u64) -> String {
    let c = ['A', 'E', 'F', 'R'];
    let mut s = String::new();
    s.push(c[((suffix as usize) >> 6) & ((1 << 2) - 1)]);
    s.push(c[((suffix as usize) >> 4) & ((1 << 2) - 1)]);
    s.push(c[((suffix as usize) >> 2) & ((1 << 2) - 1)]);
    s.push(c[(suffix as usize) & ((1 << 2) - 1)]);
    s
}

fn naive(length: usize) {
    let mut words = vec![String::new()];
    for i in 0..length {
        let mut next = vec![];
        let mut ans = vec![];
        for mut word in words.into_iter() {
            for &c in ['F', 'E', 'R', 'A'].iter() {
                word.push(c);

                let mut fare = 0;
                let mut free = 0;
                let mut area = 0;
                let mut reef = 0;
                if word.len() >= 4 {
                    for i in 0..(word.len() - 3) {
                        if &word[i..(i + 4)] == "FARE" {
                            fare += 1;
                        }
                        if &word[i..(i + 4)] == "AREA" {
                            area += 1;
                        }
                        if &word[i..(i + 4)] == "FREE" {
                            free += 1;
                        }
                        if &word[i..(i + 4)] == "REEF" {
                            reef += 1;
                        }
                    }
                }

                if fare <= 1 && free <= 1 && area <= 1 && reef <= 1 {
                    next.push(word.clone());
                }
                if fare == 1 && free == 1 && area == 1 && reef == 1 {
                    ans.push(word.clone());
                }

                word.pop();
            }
        }
        words = next;
        if i == length - 1 {
            for word in ans.into_iter() {
                println!("{}", &word[(length - 4)..length]);
            }
        }
    }
}

2019-09-24

2019/09/23

AtCoder Rust

atcoder.jp

use std::collections::{BTreeMap, BinaryHeap, VecDeque};

fn main() {
    let s = std::io::stdin();
    let mut sc = Scanner { stdin: s.lock() };
    let n: usize = sc.read();

    let mut graph = vec![vec![]; (1 << (n + 1)) - 1];
    let mut inverse = vec![vec![]; (1 << (n + 1)) - 1];
    for i in 0..((1 << n) - 1) {
        graph[i].push(i * 2 + 1);
        graph[i].push(i * 2 + 2);
        inverse[i * 2 + 1].push(i);
        inverse[i * 2 + 2].push(i);
    }
    //    eprintln!("{:?}", graph);

    let v: Vec<u64> = sc.vec(1 << n);
    let count = v.into_iter().fold(BTreeMap::new(), |mut map, v| {
        *map.entry(v).or_insert(0) += 1;
        map
    });

    let mut values = vec![0; graph.len()];
    let mut heap = BinaryHeap::new();
    heap.push((1 << n, 0));

    for (value, count) in count.into_iter().rev() {
        if heap.len() < count {
            println!("No");
            return;
        }
        let vs = (0..count).map(|_| heap.pop().unwrap()).collect::<Vec<_>>();
        for (mut size, mut cur) in vs.into_iter() {
            loop {
                values[cur] = value;
                if graph[cur].is_empty() {
                    break;
                }
                heap.push((size / 2, graph[cur][1]));
                cur = graph[cur][0];
                size /= 2;
            }
        }
    }

    println!("Yes");
}

fn get_min_leaf(root: usize, graph: &Vec<Vec<usize>>) -> usize {
    if graph[root].is_empty() {
        root
    } else {
        get_min_leaf(graph[root][0], graph)
    }
}

pub struct Scanner<R> {
    stdin: R,
}

impl<R: std::io::Read> Scanner<R> {
    pub fn read<T: std::str::FromStr>(&mut self) -> T {
        use std::io::Read;
        let buf = self
            .stdin
            .by_ref()
            .bytes()
            .map(|b| b.unwrap())
            .skip_while(|&b| b == b' ' || b == b'\n')
            .take_while(|&b| b != b' ' && b != b'\n')
            .collect::<Vec<_>>();
        unsafe { std::str::from_utf8_unchecked(&buf) }
            .parse()
            .ok()
            .expect("Parse error.")
    }
    pub fn vec<T: std::str::FromStr>(&mut self, n: usize) -> Vec<T> {
        (0..n).map(|_| self.read()).collect()
    }
    pub fn chars(&mut self) -> Vec<char> {
        self.read::<String>().chars().collect()
    }
}

atcoder.jp

fn main() {
    let s = std::io::stdin();
    let mut sc = Scanner { stdin: s.lock() };

    let n: usize = sc.read();
    let mut v: Vec<_> = (0..n)
        .map(|_| (sc.read::<f64>(), sc.read::<f64>()))
        .collect();
    v.sort_by(|&(x1, y1), &(x2, y2)| y1.atan2(x1).partial_cmp(&y2.atan2(x2)).unwrap());
    let tmp = v.clone();
    v.extend(tmp);

    let mut ans = 0.0;
    for head in 0..n {
        for tail in head..(head + n) {
            let (x, y) = (head..(tail + 1))
                .map(|i| v[i])
                .fold((0.0, 0.0), |(x, y), (vx, vy)| (x + vx, y + vy));
            update_max(&mut ans, x * x + y * y);
        }
    }
    println!("{}", ans.sqrt());
}

fn update_max<T: PartialOrd>(a: &mut T, b: T) {
    if *a < b {
        *a = b;
    }
}

pub struct Scanner<R> {
    stdin: R,
}

impl<R: std::io::Read> Scanner<R> {
    pub fn read<T: std::str::FromStr>(&mut self) -> T {
        use std::io::Read;
        let buf = self
            .stdin
            .by_ref()
            .bytes()
            .map(|b| b.unwrap())
            .skip_while(|&b| b == b' ' || b == b'\n')
            .take_while(|&b| b != b' ' && b != b'\n')
            .collect::<Vec<_>>();
        unsafe { std::str::from_utf8_unchecked(&buf) }
            .parse()
            .ok()
            .expect("Parse error.")
    }
    pub fn vec<T: std::str::FromStr>(&mut self, n: usize) -> Vec<T> {
        (0..n).map(|_| self.read()).collect()
    }
    pub fn chars(&mut self) -> Vec<char> {
        self.read::<String>().chars().collect()
    }
}

atcoder.jp

use std::collections::{BTreeMap, BTreeSet};

fn main() {
    let s = std::io::stdin();
    let mut sc = Scanner { stdin: s.lock() };
    let n = sc.read();
    let a: Vec<i64> = sc.vec(n);

    let mut ok = 0;
    let mut ng = 1e10 as i64;
    while ng - ok > 1 {
        let x = (ng + ok) / 2;
        let a = a
            .iter()
            .map(|&a| if a >= x { 1 } else { -1 })
            .collect::<Vec<i64>>();
        let mut acc = vec![0; n + 1];
        for i in 0..n {
            acc[i + 1] = acc[i] + a[i];
        }

        let map = acc
            .iter()
            .cloned()
            .collect::<BTreeSet<_>>()
            .into_iter()
            .enumerate()
            .map(|(i, v)| (v, i))
            .collect::<BTreeMap<_, _>>();
        let acc = acc
            .into_iter()
            .map(|v| *map.get(&v).unwrap())
            .collect::<Vec<_>>();
        let m = map.len();
        let mut bit = fenwick_tree::FenwickTree::new(m, 0);
        let mut negative_segments = 0;
        for i in 0..(n + 1) {
            negative_segments += i - bit.sum_one(acc[i] + 1);
            bit.add(acc[i], 1);
        }

        //        eprintln!(
        //            "x={} a={:?} acc={:?} negative_segments={}",
        //            x, a, acc, negative_segments
        //        );

        if negative_segments * 2 > (n + 1) * n / 2 {
            ng = x;
        } else {
            ok = x;
        }
    }

    println!("{}", ok);
}
pub mod fenwick_tree {
    use std::ops::{AddAssign, Sub};
    /// `FenwickTree` is a data structure that can efficiently update elements
    /// and calculate prefix sums in a table of numbers.
    /// [https://en.wikipedia.org/wiki/Fenwick_tree](https://en.wikipedia.org/wiki/Fenwick_tree)
    pub struct FenwickTree<T> {
        n: usize,
        data: Vec<T>,
        init: T,
    }

    impl<T: Copy + AddAssign + Sub<Output = T>> FenwickTree<T> {
        /// Constructs a new `FenwickTree`. The size of `FenwickTree` should be specified by `size`.
        pub fn new(size: usize, init: T) -> FenwickTree<T> {
            FenwickTree {
                n: size + 1,
                data: vec![init; size + 1],
                init: init,
            }
        }

        pub fn add(&mut self, k: usize, value: T) {
            let mut x = k;
            while x < self.n {
                self.data[x] += value;
                x |= x + 1;
            }
        }

        /// Returns a sum of range `[l, r)`
        pub fn sum(&self, l: usize, r: usize) -> T {
            self.sum_one(r) - self.sum_one(l)
        }

        /// Returns a sum of range `[0, k)`
        pub fn sum_one(&self, k: usize) -> T {
            assert!(k < self.n, "k={} n={}", k, self.n);

            let mut result = self.init;
            let mut x = k as i32 - 1;
            while x >= 0 {
                result += self.data[x as usize];
                x = (x & (x + 1)) - 1;
            }

            result
        }
    }
}

pub struct Scanner<R> {
    stdin: R,
}

impl<R: std::io::Read> Scanner<R> {
    pub fn read<T: std::str::FromStr>(&mut self) -> T {
        use std::io::Read;
        let buf = self
            .stdin
            .by_ref()
            .bytes()
            .map(|b| b.unwrap())
            .skip_while(|&b| b == b' ' || b == b'\n')
            .take_while(|&b| b != b' ' && b != b'\n')
            .collect::<Vec<_>>();
        unsafe { std::str::from_utf8_unchecked(&buf) }
            .parse()
            .ok()
            .expect("Parse error.")
    }
    pub fn vec<T: std::str::FromStr>(&mut self, n: usize) -> Vec<T> {
        (0..n).map(|_| self.read()).collect()
    }
    pub fn chars(&mut self) -> Vec<char> {
        self.read::<String>().chars().collect()
    }
}

atcoder.jp

fn main() {
    let s = std::io::stdin();
    let mut sc = Scanner { stdin: s.lock() };
    let n: usize = sc.read();
    let h: usize = sc.read();
    let w: usize = sc.read();

    let mut uf = UnionFind::new(h + w);
    let mut edges = vec![];
    for _ in 0..n {
        let row = sc.read::<usize>() - 1;
        let col = sc.read::<usize>() - 1;
        let value: u64 = sc.read();
        edges.push((value, row, col + h));
    }
    edges.sort();

    let mut ans = 0;
    for (value, row, col) in edges.into_iter().rev() {
        let x = uf.find(row);
        let y = uf.find(col);

        if x != y {
            if uf.sizes[x] + uf.sizes[y] >= uf.edge_count[x] + uf.edge_count[y] + 1 {
                uf.unite(x, y);
                ans += value;
            }
        } else if uf.sizes[x] > uf.edge_count[x] {
            uf.edge_count[x] += 1;
            ans += value;
        }
    }

    println!("{}", ans);
}

pub struct UnionFind {
    parent: Vec<usize>,
    sizes: Vec<usize>,
    size: usize,
    edge_count: Vec<usize>,
}

impl UnionFind {
    pub fn new(n: usize) -> UnionFind {
        UnionFind {
            parent: (0..n).map(|i| i).collect::<Vec<usize>>(),
            sizes: vec![1; n],
            size: n,
            edge_count: vec![0; n],
        }
    }

    pub fn find(&mut self, x: usize) -> usize {
        if x == self.parent[x] {
            x
        } else {
            let px = self.parent[x];
            self.parent[x] = self.find(px);
            self.parent[x]
        }
    }

    pub fn unite(&mut self, x: usize, y: usize) -> bool {
        let parent_x = self.find(x);
        let parent_y = self.find(y);
        if parent_x == parent_y {
            return false;
        }

        let (large, small) = if self.sizes[parent_x] < self.sizes[parent_y] {
            (parent_y, parent_x)
        } else {
            (parent_x, parent_y)
        };

        self.parent[small] = large;

        self.sizes[large] += self.sizes[small];
        self.sizes[small] = 0;

        self.edge_count[large] += self.edge_count[small] + 1;
        self.edge_count[small] = 0;

        self.size -= 1;
        return true;
    }
}

pub struct Scanner<R> {
    stdin: R,
}

impl<R: std::io::Read> Scanner<R> {
    pub fn read<T: std::str::FromStr>(&mut self) -> T {
        use std::io::Read;
        let buf = self
            .stdin
            .by_ref()
            .bytes()
            .map(|b| b.unwrap())
            .skip_while(|&b| b == b' ' || b == b'\n')
            .take_while(|&b| b != b' ' && b != b'\n')
            .collect::<Vec<_>>();
        unsafe { std::str::from_utf8_unchecked(&buf) }
            .parse()
            .ok()
            .expect("Parse error.")
    }
    pub fn vec<T: std::str::FromStr>(&mut self, n: usize) -> Vec<T> {
        (0..n).map(|_| self.read()).collect()
    }
    pub fn chars(&mut self) -> Vec<char> {
        self.read::<String>().chars().collect()
    }
}

atcoder.jp

fn main() {
    let s = std::io::stdin();
    let mut sc = Scanner { stdin: s.lock() };

    let n: usize = sc.read();
    let k: usize = sc.read();
    let power: Vec<i64> = sc.vec(n);
    if solve(&power, k) {
        println!("Yes");
    } else {
        println!("No");
    }

    // i <= hit
    // (K-(hit-i))**2
    // (K-hit+i)**2
    // i**2 + 2(K-hit)i + (K-hit)**2

    // i > hit
    // (K - (i-hit))**2
    // (K - i + hit)**2
    // (i - (hit+K))**2
    // i**2 -2(hit+K)i + (hit+K)**2
}

fn solve(power: &Vec<i64>, k: usize) -> bool {
    let n = power.len();
    let mut acc_a = vec![0; n + 1];
    let mut acc_b = vec![0; n + 1];
    let mut acc_c = vec![0; n + 1];
    for i in 0..n {
        let pi = i as i64;
        let damage = acc_a[i] * pi * pi + acc_b[i] * pi + acc_c[i];
        if damage > power[i] {
            return false;
        }

        if i + k + k - 1 > n && damage != power[i] {
            return false;
        }
        if damage < power[i] {
            let remain = power[i] - damage;
            let hit = i + k - 1;

            acc_a[i] += remain;
            acc_a[hit + k] -= remain;

            let x = hit as i64 - k as i64;
            let y = k as i64 + hit as i64;

            acc_b[i] += -remain * 2 * x;
            acc_b[hit + 1] -= -remain * 2 * x;
            acc_b[hit + 1] += -remain * 2 * y;
            acc_b[hit + k] -= -remain * 2 * y;

            acc_c[i] += remain * x * x;
            acc_c[hit + 1] -= remain * x * x;
            acc_c[hit + 1] += remain * y * y;
            acc_c[hit + k] -= remain * y * y;

            if acc_b[hit + 1].abs() > 1e15 as i64
                || acc_b[hit + k].abs() > 1e15 as i64
                || acc_c[hit + 1].abs() > 1e15 as i64
                || acc_c[hit + k].abs() > 1e15 as i64
            {
                return false;
            }
        }

        acc_a[i + 1] += acc_a[i];
        acc_b[i + 1] += acc_b[i];
        acc_c[i + 1] += acc_c[i];
    }
    true
}

pub struct Scanner<R> {
    stdin: R,
}

impl<R: std::io::Read> Scanner<R> {
    pub fn read<T: std::str::FromStr>(&mut self) -> T {
        use std::io::Read;
        let buf = self
            .stdin
            .by_ref()
            .bytes()
            .map(|b| b.unwrap())
            .skip_while(|&b| b == b' ' || b == b'\n')
            .take_while(|&b| b != b' ' && b != b'\n')
            .collect::<Vec<_>>();
        unsafe { std::str::from_utf8_unchecked(&buf) }
            .parse()
            .ok()
            .expect("Parse error.")
    }
    pub fn vec<T: std::str::FromStr>(&mut self, n: usize) -> Vec<T> {
        (0..n).map(|_| self.read()).collect()
    }
    pub fn chars(&mut self) -> Vec<char> {
        self.read::<String>().chars().collect()
    }
}

2019-09-04

AGC036 B - Do Not Duplicate

AtCoder Rust 競技プログラミング

問題

atcoder.jp

解法

s の先頭の数字と同じ数字が来たら s が空になることを考える。例えば a(0)=a(i), a(i+1)=a(j) のとき、s は i で空になり、jで再び空になる。鳩の巣原理より a0 が何周目で再び先頭になるかO(N)で求めることができる。

a0 が先頭になった場合、再び a0 が先頭になるまでに要するステップは変わらないので、最後のN*K%周期くらいを愚直にやれば良い。

コード

use std::collections::BTreeSet;

fn main() {
    let s = std::io::stdin();
    let mut sc = Scanner { stdin: s.lock() };
    let n: usize = sc.read();
    let k: usize = sc.read();
    let a: Vec<usize> = sc.vec(n);
    let max_a = a.iter().cloned().max().unwrap();
    let mut next = vec![vec![]; max_a + 1];
    let mut pos = vec![0; n];
    for (i, &a) in a.iter().enumerate() {
        pos[i] = next[a].len();
        next[a].push(i);
    }

    let mut jumped = false;
    let mut cur_t = 0;
    let mut cur_pos = 0;
    let mut prev = vec![std::usize::MAX; n];
    prev[0] = cur_t;
    loop {
        let pos = pos[cur_pos];
        if next[a[cur_pos]].len() == pos + 1 {
            if cur_t == k - 1 {
                let mut ans = vec![];
                let mut set = BTreeSet::new();
                for i in cur_pos..n {
                    if set.contains(&a[i]) {
                        while let Some(tail) = ans.pop() {
                            set.remove(&tail);
                            if tail == a[i] {
                                break;
                            }
                        }
                    } else {
                        ans.push(a[i]);
                        set.insert(a[i]);
                    }
                }

                for (i, a) in ans.into_iter().enumerate() {
                    if i > 0 {
                        print!(" ");
                    }
                    print!("{}", a);
                }
                println!();
                return;
            }
            cur_t += 1;
            cur_pos = next[a[cur_pos]][0];
        } else {
            cur_pos = next[a[cur_pos]][pos + 1]
        }
        cur_pos += 1;
        if cur_pos == n {
            if cur_t == k - 1 {
                println!();
                return;
            }
            cur_t += 1;
            cur_pos = 0;
        }
        if !jumped && prev[cur_pos] < cur_t {
            let dt = cur_t - prev[cur_pos];
            let z = (k - prev[cur_pos] - 1) / dt;
            cur_t = prev[cur_pos] + dt * z;
            assert!(cur_t < k);
            jumped = true;
        } else {
            prev[cur_pos] = cur_t;
        }
    }
}

pub struct Scanner<R> {
    stdin: R,
}

impl<R: std::io::Read> Scanner<R> {
    pub fn read<T: std::str::FromStr>(&mut self) -> T {
        use std::io::Read;
        let buf = self
            .stdin
            .by_ref()
            .bytes()
            .map(|b| b.unwrap())
            .skip_while(|&b| b == b' ' || b == b'\n')
            .take_while(|&b| b != b' ' && b != b'\n')
            .collect::<Vec<_>>();
        unsafe { std::str::from_utf8_unchecked(&buf) }
            .parse()
            .ok()
            .expect("Parse error.")
    }
    pub fn vec<T: std::str::FromStr>(&mut self, n: usize) -> Vec<T> {
        (0..n).map(|_| self.read()).collect()
    }
    pub fn chars(&mut self) -> Vec<char> {
        self.read::<String>().chars().collect()
    }
}

2019-09-02

ABC139 F - Engines

AtCoder Rust 競技プログラミング

問題

atcoder.jp

解法

答えとなるベクトル v があったとして、それを作るには v となす角が 90° 未満のベクトル集合のみを考えれば良い。よって、最終的に使われるベクトル集合は偏角ソートしたときに連続した区間になっているはずである。Nが小さいので全ての連続区間について愚直に計算する。

コード

use std::cmp::max;
use std::collections::BTreeMap;

fn main() {
    let s = std::io::stdin();
    let mut sc = Scanner { stdin: s.lock() };
    let n: usize = sc.read();
    let mut xy = vec![];
    for _ in 0..n {
        let x: i64 = sc.read();
        let y: i64 = sc.read();
        xy.push((x, y));
    }

    xy.sort_by(|&(x1, y1), &(x2, y2)| {
        let t1 = (y1 as f64).atan2(x1 as f64);
        let t2 = (y2 as f64).atan2(x2 as f64);
        t1.partial_cmp(&t2).unwrap()
    });

    let tmp = xy.clone();
    xy.extend(tmp);

    let mut ans = 0;
    for head in 0..n {
        for tail in head..(head + n) {
            let mut cur = (0, 0);
            for i in head..(tail + 1) {
                cur.0 += xy[i].0;
                cur.1 += xy[i].1;
            }
            let (x, y) = cur;
            ans = max(ans, x * x + y * y);
        }
    }
    println!("{}", (ans as f64).sqrt());
}

fn update_max<T: PartialOrd>(a: &mut T, b: T) {
    if *a < b {
        *a = b;
    }
}

fn gcd(a: i64, b: i64) -> i64 {
    if b == 0 {
        a.abs()
    } else {
        gcd(b.abs(), a.abs() % b.abs())
    }
}

pub struct Scanner<R> {
    stdin: R,
}

impl<R: std::io::Read> Scanner<R> {
    pub fn read<T: std::str::FromStr>(&mut self) -> T {
        use std::io::Read;
        let buf = self
            .stdin
            .by_ref()
            .bytes()
            .map(|b| b.unwrap())
            .skip_while(|&b| b == b' ' || b == b'\n')
            .take_while(|&b| b != b' ' && b != b'\n')
            .collect::<Vec<_>>();
        unsafe { std::str::from_utf8_unchecked(&buf) }
            .parse()
            .ok()
            .expect("Parse error.")
    }
    pub fn vec<T: std::str::FromStr>(&mut self, n: usize) -> Vec<T> {
        (0..n).map(|_| self.read()).collect()
    }
    pub fn chars(&mut self) -> Vec<char> {
        self.read::<String>().chars().collect()
    }
}

2019-09-02

ABC139 E - League

AtCoder Rust 競技プログラミング

問題

atcoder.jp

解法

順番を並び替えることを考える余地はなく、「今日試合ができる組は全員試合をする」を毎日繰り返す。

コード

use std::cmp::max;
use std::collections::VecDeque;

fn main() {
    let s = std::io::stdin();
    let mut sc = Scanner { stdin: s.lock() };

    let n: usize = sc.read();
    let a: Vec<Vec<_>> = (0..n)
        .map(|_| (1..n).map(|_| sc.read::<usize>() - 1).collect())
        .collect();
    let mut pos = vec![0; n];
    let mut q = VecDeque::new();
    for i in 0..n {
        let j = a[i][0];
        if a[j][0] == i && i < j {
            q.push_back((i, j, 1));
        }
    }

    let mut t = 0;

    while let Some((i, j, time)) = q.pop_front() {
        t = max(t, time);
        pos[i] += 1;
        if pos[i] < n - 1 {
            let j = a[i][pos[i]];
            if a[j][pos[j]] == i {
                q.push_back((i, j, time + 1));
            }
        }

        let i = j;
        pos[i] += 1;
        if pos[i] < n - 1 {
            let j = a[i][pos[i]];
            if a[j][pos[j]] == i {
                q.push_back((i, j, time + 1));
            }
        }
    }

    if pos.into_iter().all(|x| x == n - 1) {
        println!("{}", t);
    } else {
        println!("-1");
    }
}

pub struct Scanner<R> {
    stdin: R,
}

impl<R: std::io::Read> Scanner<R> {
    pub fn read<T: std::str::FromStr>(&mut self) -> T {
        use std::io::Read;
        let buf = self
            .stdin
            .by_ref()
            .bytes()
            .map(|b| b.unwrap())
            .skip_while(|&b| b == b' ' || b == b'\n')
            .take_while(|&b| b != b' ' && b != b'\n')
            .collect::<Vec<_>>();
        unsafe { std::str::from_utf8_unchecked(&buf) }
            .parse()
            .ok()
            .expect("Parse error.")
    }
    pub fn vec<T: std::str::FromStr>(&mut self, n: usize) -> Vec<T> {
        (0..n).map(|_| self.read()).collect()
    }
    pub fn chars(&mut self) -> Vec<char> {
        self.read::<String>().chars().collect()
    }
}