QOJ.ac

QOJ

IDProblemSubmitterResultTimeMemoryLanguageFile sizeSubmit timeJudge time
#320952#8213. Graffitiucup-team635#WA 0ms2308kbRust6.3kb2024-02-04 00:41:542024-02-04 00:41:54

Judging History

你现在查看的是最新测评结果

  • [2024-02-04 00:41:54]
  • 评测
  • 测评结果:WA
  • 用时:0ms
  • 内存:2308kb
  • [2024-02-04 00:41:54]
  • 提交

answer

use std::collections::*;
use std::io::Write;

type Map<K, V> = BTreeMap<K, V>;
type Set<T> = BTreeSet<T>;
type Deque<T> = VecDeque<T>;

fn main() {
    input! {
        n: usize,
        s: bytes,
        e: [(usize1, usize1); n - 1],
    }
    println!("{}", solve(s, e));
}

fn solve(s: Vec<u8>, e: Vec<(usize, usize)>) -> usize {
    let n = e.len() + 1;
    if s.len() == 1 {
        return n;
    }
    if s.len() == 2 {
        let mut res = n - 1;
        if s[0] == s[1] {
            res += n - 1;
        }
        return res;
    }
    let mut g = vec![vec![]; n];
    for (a, b) in e {
        g[a].push(b);
        g[b].push(a);
    }
    if s.windows(2).all(|s| s[0] == s[1]) {
        let ans = g
            .iter()
            .map(|g| {
                let k = g.len();
                k * (k - 1)
            })
            .sum::<usize>();
        return ans;
    }
    let root = 0;
    let mut topo = vec![root];
    for i in 0..n {
        let v = topo[i];
        for u in g[v].clone() {
            g[u].retain(|p| *p != v);
            topo.push(u);
        }
    }
    if s[0] != s[1] && s[1] != s[2] {
        // 根がs[1] である、親がs[1]である = 最大値
        let mut dp = vec![[[0i64; 2]; 2]; n];
        for &v in topo.iter().rev() {
            let child = g[v].iter().map(|u| dp[*u]).collect::<Vec<_>>();
            let mut val = [[0; 2]; 2];
            for (now, val) in val.iter_mut().enumerate() {
                for (parent, val) in val.iter_mut().enumerate() {
                    let c = child
                        .iter()
                        .map(|c| (c[0][now], c[1][now]))
                        .collect::<Vec<_>>();
                    let mut sum = c.iter().map(|p| p.0).sum::<i64>();
                    let mut c = c.iter().map(|p| p.1 - p.0).collect::<Vec<_>>();
                    c.sort();
                    let mut k = c.len() as i64;
                    if parent == 0 {
                        k += 1;
                    }
                    let now = now as i64;
                    if s[0] == s[2] {
                        val.chmax(sum + now * k * (k - 1));
                    } else {
                        val.chmax(sum + now * (k / 2) * (k - k / 2));
                    }
                    for _ in 0..c.len() {
                        sum += c.pop().unwrap();
                        k -= 1;
                        if s[0] == s[2] {
                            val.chmax(sum + now * k * (k - 1));
                        } else {
                            val.chmax(sum + now * (k / 2) * (k - k / 2));
                        }
                    }
                }
            }
            dp[v] = val;
        }
        let mut ans = dp[root][0][1].max(dp[root][1][1]);
        if s[0] == s[2] {
            ans *= 2;
        }
        return ans as usize;
    }
    // 001 と言う文字列を考える
    // 根、親 = 最大値
    let mut dp = vec![[[0i64; 2]; 2]; n];
    for &v in topo.iter().rev() {
        let child = g[v].iter().map(|u| dp[*u]).collect::<Vec<_>>();
        let mut val = [[0; 2]; 2];
        for (now, val) in val.iter_mut().enumerate() {
            for (parent, val) in val.iter_mut().enumerate() {
                let c = child
                    .iter()
                    .map(|c| (c[0][now], c[1][now]))
                    .collect::<Vec<_>>();
                let mut sum = c.iter().map(|p| p.0).sum::<i64>();
                let mut c = c.iter().map(|p| p.1 - p.0).collect::<Vec<_>>();
                c.sort();
                let mut zero = c.len() as i64;
                let mut one = 0;
                if v != root {
                    if parent == 0 {
                        zero += 1;
                    } else {
                        one += 1;
                    }
                }
                let now = (now ^ 1) as i64;
                val.chmax(sum + now * zero * one);
                for _ in 0..c.len() {
                    sum += c.pop().unwrap();
                    zero -= 1;
                    one += 1;
                    val.chmax(sum + now * zero * one);
                }
            }
        }
        dp[v] = val;
    }
    let ans = dp[root][0][0].max(dp[root][1][0]);
    ans as usize
}

// ---------- begin input macro ----------
// reference: https://qiita.com/tanakh/items/0ba42c7ca36cd29d0ac8
#[macro_export]
macro_rules! input {
    (source = $s:expr, $($r:tt)*) => {
        let mut iter = $s.split_whitespace();
        input_inner!{iter, $($r)*}
    };
    ($($r:tt)*) => {
        let s = {
            use std::io::Read;
            let mut s = String::new();
            std::io::stdin().read_to_string(&mut s).unwrap();
            s
        };
        let mut iter = s.split_whitespace();
        input_inner!{iter, $($r)*}
    };
}

#[macro_export]
macro_rules! input_inner {
    ($iter:expr) => {};
    ($iter:expr, ) => {};
    ($iter:expr, $var:ident : $t:tt $($r:tt)*) => {
        let $var = read_value!($iter, $t);
        input_inner!{$iter $($r)*}
    };
}

#[macro_export]
macro_rules! read_value {
    ($iter:expr, ( $($t:tt),* )) => {
        ( $(read_value!($iter, $t)),* )
    };
    ($iter:expr, [ $t:tt ; $len:expr ]) => {
        (0..$len).map(|_| read_value!($iter, $t)).collect::<Vec<_>>()
    };
    ($iter:expr, chars) => {
        read_value!($iter, String).chars().collect::<Vec<char>>()
    };
    ($iter:expr, bytes) => {
        read_value!($iter, String).bytes().collect::<Vec<u8>>()
    };
    ($iter:expr, usize1) => {
        read_value!($iter, usize) - 1
    };
    ($iter:expr, $t:ty) => {
        $iter.next().unwrap().parse::<$t>().expect("Parse error")
    };
}
// ---------- end input macro ----------
// ---------- begin chmin, chmax ----------
pub trait ChangeMinMax {
    fn chmin(&mut self, x: Self) -> bool;
    fn chmax(&mut self, x: Self) -> bool;
}

impl<T: PartialOrd> ChangeMinMax for T {
    fn chmin(&mut self, x: Self) -> bool {
        *self > x && {
            *self = x;
            true
        }
    }
    fn chmax(&mut self, x: Self) -> bool {
        *self < x && {
            *self = x;
            true
        }
    }
}
// ---------- end chmin, chmax ----------

Details

Tip: Click on the bar to expand more detailed information

Test #1:

score: 100
Accepted
time: 0ms
memory: 2084kb

input:

1
a

output:

1

result:

ok 1 number(s): "1"

Test #2:

score: 0
Accepted
time: 0ms
memory: 2108kb

input:

3
orz
1 2
2 3

output:

1

result:

ok 1 number(s): "1"

Test #3:

score: 0
Accepted
time: 0ms
memory: 2120kb

input:

2
ab
1 2

output:

1

result:

ok 1 number(s): "1"

Test #4:

score: -100
Wrong Answer
time: 0ms
memory: 2308kb

input:

5
bob
3 2
5 1
1 4
2 4

output:

8

result:

wrong answer 1st numbers differ - expected: '4', found: '8'