QOJ.ac
QOJ
ID | Problem | Submitter | Result | Time | Memory | Language | File size | Submit time | Judge time |
---|---|---|---|---|---|---|---|---|---|
#320952 | #8213. Graffiti | ucup-team635# | WA | 0ms | 2308kb | Rust | 6.3kb | 2024-02-04 00:41:54 | 2024-02-04 00:41:54 |
Judging History
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'