// https://judge.eluminatis-of-lu.com/contest/676ffd92569fb90008aac7da/1157
use crate::algo_lib::collections::treap::payload::{OrdPayload, Payload};
use crate::algo_lib::collections::treap::Tree;
use crate::algo_lib::collections::vec_ext::inc_dec::IncDec;
use crate::algo_lib::graph::edges::edge_trait::EdgeTrait;
use crate::algo_lib::graph::Graph;
use crate::algo_lib::io::input::Input;
use crate::algo_lib::io::output::Output;
use crate::algo_lib::misc::recursive_function::{Callable2, RecursiveFunction2};
use crate::algo_lib::misc::test_type::TaskType;
use crate::algo_lib::misc::test_type::TestType;
type PreCalc = ();
fn solve(input: &mut Input, out: &mut Output, _test_case: usize, _data: &mut PreCalc) {
let n = input.read_size();
let a = input.read_size_vec(n).dec();
let edges = input.read_size_pair_vec(n - 1).dec();
struct Node(usize);
impl Payload for Node {}
impl OrdPayload for Node {
type Key = usize;
fn key(&self) -> &Self::Key {
&self.0
}
fn union(a: Self, _b: Self) -> Self {
a
}
}
let mut ans = vec![0; n];
let graph = Graph::from_biedges(n, &edges);
let mut dfs = RecursiveFunction2::new(|
f,
vert: usize,
prev: usize,
| -> (Tree<Node>, usize) {
let mut treap = Tree::new();
treap.insert(Node(a[vert]));
let mut size = 1;
for e in &graph[vert] {
if e.to() == prev {
continue;
}
let (call, call_sz) = f.call(e.to(), vert);
treap = Tree::union(treap, call);
size += call_sz;
}
ans[vert] = size - treap.range(..&size).size();
(treap, size)
});
dfs.call(0, n);
let q = input.read_size();
for _ in 0..q {
let x = input.read_size() - 1;
out.print_line(ans[x]);
}
}
pub static TEST_TYPE: TestType = TestType::MultiNumber;
pub static TASK_TYPE: TaskType = TaskType::Classic;
pub(crate) fn run(mut input: Input, mut output: Output) -> bool {
let mut pre_calc = ();
match TEST_TYPE {
TestType::Single => solve(&mut input, &mut output, 1, &mut pre_calc),
TestType::MultiNumber => {
let t = input.read();
for i in 1..=t {
solve(&mut input, &mut output, i, &mut pre_calc);
}
}
TestType::MultiEof => {
let mut i = 1;
while input.peek().is_some() {
solve(&mut input, &mut output, i, &mut pre_calc);
i += 1;
}
}
}
output.flush();
match TASK_TYPE {
TaskType::Classic => input.is_empty(),
TaskType::Interactive => true,
}
}
fn main() {
let mut sin = std::io::stdin();
let input = crate::algo_lib::io::input::Input::new(&mut sin);
let mut stdout = std::io::stdout();
let output = crate::algo_lib::io::output::Output::new(&mut stdout);
run(input, output);
}
pub mod algo_lib {
pub mod collections {
pub mod dsu {
use crate::algo_lib::collections::slice_ext::bounds::Bounds;
use crate::algo_lib::collections::slice_ext::indices::Indices;
use std::cell::Cell;
#[derive(Clone)]
pub struct DSU {
id: Vec<Cell<i32>>,
count: usize,
}
impl DSU {
pub fn new(n: usize) -> Self {
Self {
id: vec![Cell::new(- 1); n],
count: n,
}
}
pub fn size(&self, i: usize) -> usize {
(-self.id[self.find(i)].get()) as usize
}
#[allow(clippy::len_without_is_empty)]
pub fn len(&self) -> usize {
self.id.len()
}
pub fn iter(&self) -> impl Iterator<Item = usize> + '_ {
self.id
.iter()
.enumerate()
.filter_map(|(i, id)| if id.get() < 0 { Some(i) } else { None })
}
pub fn set_count(&self) -> usize {
self.count
}
pub fn union(&mut self, mut a: usize, mut b: usize) -> bool {
a = self.find(a);
b = self.find(b);
if a == b {
false
} else {
self.id[a].replace(self.id[a].get() + self.id[b].get());
self.id[b].replace(a as i32);
self.count -= 1;
true
}
}
pub fn find(&self, i: usize) -> usize {
if self.id[i].get() >= 0 {
let res = self.find(self.id[i].get() as usize);
self.id[i].replace(res as i32);
res
} else {
i
}
}
pub fn clear(&mut self) {
self.count = self.id.len();
self.id.fill(Cell::new(-1));
}
pub fn parts(&self) -> Vec<Vec<usize>> {
let roots: Vec<_> = self.iter().collect();
let mut res = vec![Vec::new(); roots.len()];
for i in self.id.indices() {
res[roots.as_slice().bin_search(&self.find(i)).unwrap()].push(i);
}
res
}
}
}
pub mod slice_ext {
pub mod bounds {
pub trait Bounds<T: PartialOrd> {
fn lower_bound(&self, el: &T) -> usize;
fn upper_bound(&self, el: &T) -> usize;
fn bin_search(&self, el: &T) -> Option<usize>;
fn more(&self, el: &T) -> usize;
fn more_or_eq(&self, el: &T) -> usize;
fn less(&self, el: &T) -> usize {
self.lower_bound(el)
}
fn less_or_eq(&self, el: &T) -> usize {
self.upper_bound(el)
}
}
impl<T: PartialOrd> Bounds<T> for [T] {
fn lower_bound(&self, el: &T) -> usize {
let mut left = 0;
let mut right = self.len();
while left < right {
let mid = left + ((right - left) >> 1);
if &self[mid] < el {
left = mid + 1;
} else {
right = mid;
}
}
left
}
fn upper_bound(&self, el: &T) -> usize {
let mut left = 0;
let mut right = self.len();
while left < right {
let mid = left + ((right - left) >> 1);
if &self[mid] <= el {
left = mid + 1;
} else {
right = mid;
}
}
left
}
fn bin_search(&self, el: &T) -> Option<usize> {
let at = self.lower_bound(el);
if at == self.len() || &self[at] != el { None } else { Some(at) }
}
fn more(&self, el: &T) -> usize {
self.len() - self.upper_bound(el)
}
fn more_or_eq(&self, el: &T) -> usize {
self.len() - self.lower_bound(el)
}
}
}
pub mod indices {
use std::ops::Range;
pub trait Indices {
fn indices(&self) -> Range<usize>;
}
impl<T> Indices for [T] {
fn indices(&self) -> Range<usize> {
0..self.len()
}
}
}
}
pub mod treap {
use crate::algo_lib::collections::treap::payload::{OrdPayload, Payload, Pushable};
use crate::algo_lib::misc::direction::Direction;
use crate::algo_lib::misc::extensions::replace_with::ReplaceWith;
use crate::algo_lib::misc::random::{gen, gen_range};
use std::collections::Bound;
use std::marker::{PhantomData, PhantomPinned};
use std::mem::{forget, replace, swap, take, MaybeUninit};
use std::ops::{Deref, DerefMut, RangeBounds};
use std::ptr::NonNull;
pub struct Content<P> {
payload: P,
parent: NodeLink<Node<P>>,
left: NodeLink<Node<P>>,
right: NodeLink<Node<P>>,
}
impl<P: Payload> Content<P> {
fn push_down(&mut self) {
if P::NEED_PUSH_DOWN && self.payload.need_push_down() {
if self.left.size != 0 {
self.left.deref_mut().payload_mut().push_delta(&self.payload);
}
if self.right.size != 0 {
self.right.deref_mut().payload_mut().push_delta(&self.payload);
}
self.payload.reset_delta();
}
}
fn update(&mut self) {
if P::NEED_UPDATE {
self.payload.update(self.left.payload(), self.right.payload());
}
}
}
pub struct Node<P> {
priority: u64,
size: u32,
reversed: bool,
content: Option<Content<P>>,
_phantom_pinned: PhantomPinned,
}
impl<P> Node<P> {
const NULL_NODE: Self = Self {
priority: 0,
size: 0,
reversed: false,
content: None,
_phantom_pinned: PhantomPinned,
};
fn new(payload: P) -> NodeLink<Node<P>> {
NodeLink::new(Node {
priority: gen(),
size: 1,
reversed: false,
content: Some(Content {
payload,
parent: NodeLink::default(),
left: NodeLink::default(),
right: NodeLink::default(),
}),
_phantom_pinned: PhantomPinned,
})
}
fn payload(&self) -> Option<&P> {
self.content.as_ref().map(|c| &c.payload)
}
fn payload_mut(&mut self) -> &mut P {
unsafe { &mut self.content.as_mut().unwrap_unchecked().payload }
}
fn reverse(&mut self) {
if let Some(content) = &mut self.content {
self.reversed = !self.reversed;
swap(&mut content.left, &mut content.right);
}
}
}
impl<P: Payload> Node<P> {
fn update(&mut self) {
self.size = 1 + self.left.size + self.right.size;
self.deref_mut().update();
}
fn push_down(&mut self) {
if self.reversed {
self.left.reverse();
self.right.reverse();
self.reversed = false;
}
self.deref_mut().push_down();
}
fn detach(&mut self) {
if self.size != 0 {
self.parent = NodeLink::default();
}
}
fn detach_left(&mut self) -> NodeLink<Node<P>> {
self.push_down();
let mut left = take(&mut self.left);
left.detach();
left
}
fn detach_right(&mut self) -> NodeLink<Node<P>> {
self.push_down();
let mut right = take(&mut self.right);
right.detach();
right
}
fn attach_left(&mut self, left: NodeLink<Node<P>>) {
assert_eq!(self.left.size, 0);
if left.size != 0 {
self.left = left;
self.left.deref_mut().parent = NodeLink::new_ref(self);
}
self.update();
}
fn attach_right(&mut self, right: NodeLink<Node<P>>) {
assert_eq!(self.right.size, 0);
if right.size != 0 {
self.right = right;
self.right.deref_mut().parent = NodeLink::new_ref(self);
}
self.update();
}
fn heapify(&mut self) {
if self.left.size != 0 {
self.left.heapify();
if self.left.priority > self.priority {
let left_priority = self.left.priority;
self.left.priority = self.priority;
self.priority = left_priority;
}
}
if self.right.size != 0 {
self.right.heapify();
if self.right.priority > self.priority {
let right_priority = self.right.priority;
self.right.priority = self.priority;
self.priority = right_priority;
}
}
}
fn gen(n: usize, f: impl FnMut(usize) -> P) -> NodeLink<Node<P>> {
let mut res = Self::build(f, 0, n).0;
res.heapify();
res
}
fn build<F: FnMut(usize) -> P>(
mut f: F,
from: usize,
to: usize,
) -> (NodeLink<Node<P>>, F) {
if from == to {
(NodeLink::default(), f)
} else {
let mid = gen_range(from..to);
let mut node = Node::new(f(mid));
let (left, f) = Self::build(f, from, mid);
let (right, f) = Self::build(f, mid + 1, to);
node.attach_left(left);
node.attach_right(right);
(node, f)
}
}
fn push<D>(&mut self, delta: D)
where
P: Pushable<D>,
{
self.payload_mut().push(delta);
}
fn refs(&mut self, res: &mut Vec<NodeId<P>>) {
if self.size != 0 {
self.left.refs(res);
res.push(NodeId(NodeLink::new_ref(self)));
self.right.refs(res);
}
}
fn first(&mut self) -> &Node<P> {
if self.size == 0 {
&Self::NULL_NODE
} else {
self.push_down();
if self.left.size != 0 { self.left.first() } else { self }
}
}
fn last(&mut self) -> &Node<P> {
if self.size == 0 {
&Self::NULL_NODE
} else {
self.push_down();
if self.right.size != 0 { self.right.last() } else { self }
}
}
fn binary_search<F: FnMut(&P, Option<&P>, Option<&P>) -> Option<Direction>>(
&mut self,
mut f: F,
) {
if self.size != 0 {
self.push_down();
let direction = f(&self.payload, self.left.payload(), self.right.payload());
match direction {
Some(Direction::Left) => self.left.binary_search(f),
Some(Direction::Right) => self.right.binary_search(f),
None => {}
}
}
}
fn size(&self) -> usize {
self.size as usize
}
fn binary_search_size<F: FnMut(usize, usize) -> Option<Direction>>(
&mut self,
mut f: F,
) {
if self.size != 0 {
self.push_down();
let direction = f(self.left.size(), self.right.size());
match direction {
Some(Direction::Left) => self.left.binary_search_size(f),
Some(Direction::Right) => self.right.binary_search_size(f),
None => {}
}
}
}
}
impl<P> Deref for Node<P> {
type Target = Content<P>;
fn deref(&self) -> &Self::Target {
unsafe { self.content.as_ref().unwrap_unchecked() }
}
}
impl<P> DerefMut for Node<P> {
fn deref_mut(&mut self) -> &mut Self::Target {
unsafe { self.content.as_mut().unwrap_unchecked() }
}
}
pub struct NodeId<P>(NodeLink<Node<P>>);
pub struct NodeLink<P> {
link: NonNull<P>,
}
impl<P> NodeLink<Node<P>> {
fn new(node: Node<P>) -> Self {
let mut pinned = Box::pin(node);
let res = NodeLink {
link: unsafe { NonNull::from(pinned.as_mut().get_unchecked_mut()) },
};
forget(pinned);
res
}
fn new_ref(node: &Node<P>) -> Self {
NodeLink {
link: NonNull::from(node),
}
}
fn into_payload(mut self) -> P {
assert_eq!(self.left.size, 0);
assert_eq!(self.right.size, 0);
assert_eq!(self.parent.size, 0);
self.content.take().unwrap().payload
}
}
impl<P: Payload> NodeLink<Node<P>> {
fn merge(mut left: Self, mut right: Self) -> Self {
match (left.size, right.size) {
(0, _) => right,
(_, 0) => left,
_ => {
if left.priority > right.priority {
let left_right = left.detach_right();
left.attach_right(Self::merge(left_right, right));
left
} else {
let right_left = right.detach_left();
right.attach_left(Self::merge(left, right_left));
right
}
}
}
}
fn split_by<F: FnMut(&P, Option<&P>, Option<&P>) -> Direction>(
mut self,
mut f: F,
) -> (Self, Self) {
if self.size != 0 {
let direction = f(&self.payload, self.left.payload(), self.right.payload());
match direction {
Direction::Left => {
let (left, right) = self.detach_left().split_by(f);
self.attach_left(right);
(left, self)
}
Direction::Right => {
let (left, right) = self.detach_right().split_by(f);
self.attach_right(left);
(self, right)
}
}
} else {
(NodeLink::default(), NodeLink::default())
}
}
fn push_from_up(&self, directions: &mut Vec<Direction>) -> NodeLink<Node<P>> {
if self.parent.size != 0 {
if self.parent.left == *self {
directions.push(Direction::Left);
} else if self.parent.right == *self {
directions.push(Direction::Right);
} else {
unreachable!();
}
self.parent.push_from_up(directions)
} else {
NodeLink { link: self.link }
}
}
fn raise(self, link: &NodeLink<Node<P>>) -> (Self, Self, Self) {
assert!(link.content.is_some());
let mut directions = Vec::new();
let expected_parent = link.push_from_up(&mut directions);
assert!(expected_parent == self);
self.split_by_dir(directions)
}
fn split_by_dir(mut self, mut directions: Vec<Direction>) -> (Self, Self, Self) {
if let Some(dir) = directions.pop() {
match dir {
Direction::Left => {
let (left, mid, right) = self.detach_left().split_by_dir(directions);
self.attach_left(right);
(left, mid, self)
}
Direction::Right => {
let (left, mid, right) = self
.detach_right()
.split_by_dir(directions);
self.attach_right(left);
(self, mid, right)
}
}
} else {
let left = self.detach_left();
let right = self.detach_right();
self.update();
(left, self, right)
}
}
fn split_by_size<F: FnMut(usize, usize) -> Direction>(
mut self,
mut f: F,
) -> (Self, Self) {
if self.size != 0 {
let direction = f(self.left.size(), self.right.size());
match direction {
Direction::Left => {
let (left, right) = self.detach_left().split_by_size(f);
self.attach_left(right);
(left, self)
}
Direction::Right => {
let (left, right) = self.detach_right().split_by_size(f);
self.attach_right(left);
(self, right)
}
}
} else {
(Self::default(), Self::default())
}
}
fn split_at(self, mut at: usize) -> (Self, Self) {
self.split_by_size(|left_size, _| {
if at <= left_size {
Direction::Left
} else {
at -= left_size + 1;
Direction::Right
}
})
}
}
impl<P: OrdPayload> NodeLink<Node<P>> {
fn split(self, key: &P::Key) -> (Self, Self) {
self.split_by(|payload, _left, _right| {
if payload.key() < key { Direction::Right } else { Direction::Left }
})
}
fn split_inclusive(self, key: &P::Key) -> (Self, Self) {
self.split_by(|payload, _left, _right| {
if payload.key() <= key { Direction::Right } else { Direction::Left }
})
}
fn union(mut a: Self, mut b: Self) -> Self {
match (a.size, b.size) {
(0, _) => b,
(_, 0) => a,
_ => {
if a.priority < b.priority {
swap(&mut a, &mut b);
}
let (b_left, b_right) = b.split(a.payload.key());
let (same, b_right) = b_right.split_inclusive(a.payload.key());
let (left, right) = if same.size != 0 {
let left = a.detach_left();
let right = a.detach_right();
a = Node::new(P::union(a.into_payload(), same.into_payload()));
(left, right)
} else {
(a.detach_left(), a.detach_right())
};
let left = Self::union(left, b_left);
let right = Self::union(right, b_right);
a.attach_left(left);
a.attach_right(right);
a
}
}
}
}
impl<P> Clone for NodeLink<P> {
fn clone(&self) -> Self {
NodeLink { link: self.link }
}
}
impl<P> PartialEq for NodeLink<P> {
fn eq(&self, other: &Self) -> bool {
self.link == other.link
}
}
impl<P> Eq for NodeLink<P> {}
impl<P> Deref for NodeLink<P> {
type Target = P;
fn deref(&self) -> &Self::Target {
unsafe { self.link.as_ref() }
}
}
impl<P> DerefMut for NodeLink<P> {
fn deref_mut(&mut self) -> &mut Self::Target {
unsafe { self.link.as_mut() }
}
}
impl<P> Default for NodeLink<Node<P>> {
fn default() -> Self {
NodeLink {
link: unsafe {
NonNull::new_unchecked(
&Node::NULL_NODE as *const Node<P> as *mut Node<P>,
)
},
}
}
}
pub enum Tree<P> {
Whole { root: NodeLink<Node<P>> },
Split { left: NodeLink<Node<P>>, mid: Box<Tree<P>>, right: NodeLink<Node<P>> },
}
impl<P: Payload> Default for Tree<P> {
fn default() -> Self {
Self::new()
}
}
impl<P: Payload> Tree<P> {
pub fn new() -> Self {
Tree::Whole {
root: NodeLink::default(),
}
}
pub fn gen(n: usize, f: impl FnMut(usize) -> P) -> Self {
Self::gen_impl(n, f)
}
fn single(p: P) -> Self {
Tree::Whole { root: Node::new(p) }
}
fn gen_impl(n: usize, f: impl FnMut(usize) -> P) -> Self {
Tree::Whole {
root: Node::gen(n, f),
}
}
fn into_node(mut self) -> NodeLink<Node<P>> {
self.rebuild();
match self {
Tree::Whole { root } => root,
_ => unreachable!(),
}
}
fn as_node(&mut self) -> &mut NodeLink<Node<P>> {
match self {
Tree::Whole { root } => root,
_ => unreachable!(),
}
}
fn rebuild(&mut self) -> &mut NodeLink<Node<P>> {
self.replace_with(|self_| {
if let Tree::Split { left, mid, right } = self_ {
Tree::Whole {
root: NodeLink::merge(left, NodeLink::merge(mid.into_node(), right)),
}
} else {
self_
}
});
self.as_node()
}
fn do_split(
mut self,
f: impl FnOnce(NodeLink<Node<P>>) -> (NodeLink<Node<P>>, NodeLink<Node<P>>),
) -> (Self, Self) {
let mut right = MaybeUninit::uninit();
self.replace_with(|self_| {
let root = self_.into_node();
let (left, right_) = f(root);
right.write(Tree::Whole { root: right_ });
Tree::Whole { root: left }
});
(self, unsafe { right.assume_init() })
}
fn do_split_three(
&mut self,
f: impl FnOnce(
NodeLink<Node<P>>,
) -> (NodeLink<Node<P>>, NodeLink<Node<P>>, NodeLink<Node<P>>),
) -> &mut Self {
self.replace_with(|self_| {
let root = self_.into_node();
let (left, mid, right) = f(root);
Self::Split {
left,
mid: Box::new(Self::Whole { root: mid }),
right,
}
});
self.as_mid()
}
fn as_mid(&mut self) -> &mut Self {
match self {
Tree::Split { mid, .. } => mid,
_ => unreachable!(),
}
}
pub fn push_front(&mut self, other: Self) -> &mut Self {
self.replace_with(|root| Tree::Whole {
root: NodeLink::merge(other.into_node(), root.into_node()),
});
self
}
pub fn push_back(&mut self, other: Self) -> &mut Self {
self.replace_with(|root| Tree::Whole {
root: NodeLink::merge(root.into_node(), other.into_node()),
});
self
}
pub fn detach(&mut self) -> Self {
match self {
Tree::Whole { root } => Tree::Whole { root: take(root) },
Tree::Split { mid, .. } => replace(mid, Tree::new()),
}
}
pub fn binary_search(
&mut self,
f: impl FnMut(&P, Option<&P>, Option<&P>) -> Option<Direction>,
) {
self.rebuild().binary_search(f);
}
pub fn push<D>(&mut self, delta: D)
where
P: Pushable<D>,
{
self.rebuild().push(delta);
}
pub fn merge(left: Self, right: Self) -> Self {
match left {
Tree::Whole { root: left_root } => {
Tree::Split {
left: left_root,
mid: Box::new(right),
right: NodeLink::default(),
}
}
Tree::Split { left, mid, right: left_right } => {
Tree::Split {
left,
mid,
right: NodeLink::merge(left_right, right.into_node()),
}
}
}
}
pub fn merge_three(left: Self, mid: Self, right: Self) -> Self {
Self::Split {
left: left.into_node(),
mid: Box::new(mid),
right: right.into_node(),
}
}
pub fn iter(&mut self) -> Iter<'_, P> {
Iter::new(self.rebuild())
}
pub fn first(&mut self) -> Option<&P> {
self.rebuild().first().payload()
}
pub fn last(&mut self) -> Option<&P> {
self.rebuild().last().payload()
}
pub fn payload(&mut self) -> Option<&P> {
self.rebuild().payload()
}
pub fn add_back(&mut self, payload: P) -> NodeId<P> {
let mut res = MaybeUninit::uninit();
self.replace_with(|root| {
let mut new_node = Self::single(payload);
res.write(NodeId(new_node.as_node().clone()));
Self::merge(root, new_node)
});
unsafe { res.assume_init() }
}
pub fn add_front(&mut self, payload: P) -> NodeId<P> {
let mut res = MaybeUninit::uninit();
self.replace_with(|root| {
let mut new_node = Self::single(payload);
res.write(NodeId(new_node.as_node().clone()));
Self::merge(new_node, root)
});
unsafe { res.assume_init() }
}
pub fn refs(&mut self) -> Vec<NodeId<P>> {
let mut res = Vec::new();
self.rebuild().refs(&mut res);
res
}
pub fn is_empty(&self) -> bool {
match self {
Tree::Whole { root } => root.size == 0,
Tree::Split { left, mid, right } => {
left.size == 0 && right.size == 0 && mid.is_empty()
}
}
}
pub fn raise(&mut self, node_ref: &NodeId<P>) -> &mut Self {
self.do_split_three(|node| node.raise(&node_ref.0))
}
pub fn into_payload(self) -> P {
self.into_node().into_payload()
}
pub fn index_ref(&mut self, node_ref: &NodeId<P>) -> usize {
self.raise(node_ref);
match self {
Tree::Whole { .. } => unreachable!(),
Tree::Split { left, .. } => left.size(),
}
}
pub fn size(&self) -> usize {
match self {
Tree::Whole { root } => root.size(),
Tree::Split { left, mid, right } => left.size() + mid.size() + right.size(),
}
}
pub fn split_at(self, at: usize) -> (Self, Self) {
self.do_split(|root| root.split_at(at))
}
pub fn binary_search_size(
&mut self,
f: impl FnMut(usize, usize) -> Option<Direction>,
) {
self.rebuild().binary_search_size(f);
}
pub fn range_index(&mut self, bounds: impl RangeBounds<usize>) -> &mut Self {
self.do_split_three(|root| {
let size = root.size();
let start = match bounds.start_bound() {
Bound::Included(&s) => s,
Bound::Excluded(&s) => s + 1,
Bound::Unbounded => 0,
};
let end = match bounds.end_bound() {
Bound::Included(&e) => e + 1,
Bound::Excluded(&e) => e,
Bound::Unbounded => size,
};
assert!(start <= end);
let (left, mid_right) = root.split_at(start);
let (mid, right) = mid_right.split_at(end.max(start) - start);
(left, mid, right)
})
}
pub fn get_at(&mut self, at: usize) -> &mut Self {
self.range_index(at..=at)
}
pub fn reverse(&mut self) {
self.rebuild().reverse();
}
}
pub struct Iter<'a, P> {
stack: Vec<*mut NodeLink<Node<P>>>,
_marker: PhantomData<&'a P>,
}
impl<'a, P: Payload> Iter<'a, P> {
fn new(root: &'a mut NodeLink<Node<P>>) -> Self {
let mut res = Self {
stack: Vec::new(),
_marker: PhantomData,
};
res.add_left(root);
res
}
fn add_left(&mut self, mut node: &mut NodeLink<Node<P>>) {
while node.size != 0 {
node.push_down();
self.stack.push(node);
node = &mut node.left;
}
}
}
impl<'a, P: Payload> Iterator for Iter<'a, P> {
type Item = &'a P;
fn next(&mut self) -> Option<Self::Item> {
let node = unsafe { &mut *self.stack.pop()? };
self.add_left(&mut node.right);
Some(&node.payload)
}
}
impl<P: OrdPayload> Tree<P> {
pub fn range<'s: 'r, 'r>(
&'s mut self,
bounds: impl RangeBounds<&'r P::Key>,
) -> &mut Self {
self.do_split_three(|root| {
let (left, mid_right) = match bounds.start_bound() {
Bound::Included(key) => root.split(key),
Bound::Excluded(key) => root.split_inclusive(key),
Bound::Unbounded => (NodeLink::default(), root),
};
let (mid, right) = match bounds.end_bound() {
Bound::Included(key) => mid_right.split_inclusive(key),
Bound::Excluded(key) => mid_right.split(key),
Bound::Unbounded => (mid_right, NodeLink::default()),
};
(left, mid, right)
})
}
pub fn insert(&mut self, p: P) -> Option<P> {
let mid = self.range(&p.key()..=&p.key());
let mut res = None;
mid.replace_with(|mid| {
if mid.size() != 0 {
res = Some(mid.into_node().into_payload());
}
Tree::single(p)
});
res
}
pub fn remove(&mut self, key: &P::Key) -> Option<P> {
let mid = self.range(key..=key);
let mut res = None;
mid.replace_with(|mid| {
if mid.size() != 0 {
res = Some(mid.into_node().into_payload());
}
Self::new()
});
res
}
pub fn split(self, key: &P::Key) -> (Self, Self) {
self.do_split(|root| root.split(key))
}
pub fn split_inclusive(self, key: &P::Key) -> (Self, Self) {
self.do_split(|root| root.split_inclusive(key))
}
pub fn get(&mut self, key: &P::Key) -> Option<&P> {
self.range(key..=key).payload()
}
pub fn floor(&mut self, key: &P::Key) -> Option<&P> {
self.range(..=key).last()
}
pub fn ceil(&mut self, key: &P::Key) -> Option<&P> {
self.range(key..).first()
}
pub fn prev(&mut self, key: &P::Key) -> Option<&P> {
self.range(..key).last()
}
pub fn next(&mut self, key: &P::Key) -> Option<&P> {
self.range((Bound::Excluded(key), Bound::Unbounded)).first()
}
pub fn union(a: Self, b: Self) -> Self {
let a = a.into_node();
let b = b.into_node();
Self::Whole {
root: NodeLink::union(a, b),
}
}
pub fn index(&mut self, key: &P::Key) -> Option<usize> {
match self.range(key..=key) {
Tree::Split { left, mid, .. } => mid.payload().map(|_| left.size()),
_ => unreachable!(),
}
}
}
pub mod payload {
#[allow(unused_variables)]
pub trait Payload: Sized {
const NEED_UPDATE: bool = false;
const NEED_PUSH_DOWN: bool = false;
fn reset_delta(&mut self) {
unimplemented!()
}
fn update(&mut self, left: Option<&Self>, right: Option<&Self>) {
unimplemented!()
}
fn push_delta(&mut self, delta: &Self) {
unimplemented!()
}
fn need_push_down(&self) -> bool {
unimplemented!()
}
}
#[allow(unused_variables)]
pub trait OrdPayload: Payload {
type Key: Ord;
fn key(&self) -> &Self::Key;
fn union(a: Self, b: Self) -> Self {
unimplemented!()
}
}
pub trait Pushable<Delta>: Payload {
fn push(&mut self, delta: Delta);
}
impl<P: Payload> Pushable<&P> for P {
#[inline]
fn push(&mut self, delta: &P) {
self.push_delta(delta);
}
}
impl<P: Payload> Pushable<P> for P {
#[inline]
fn push(&mut self, delta: P) {
*self = delta;
}
}
}
}
pub mod vec_ext {
pub mod default {
pub fn default_vec<T: Default>(len: usize) -> Vec<T> {
let mut v = Vec::with_capacity(len);
for _ in 0..len {
v.push(T::default());
}
v
}
}
pub mod inc_dec {
use crate::algo_lib::numbers::num_traits::algebra::{AdditionMonoidWithSub, One};
pub trait IncDec {
#[must_use]
fn inc(self) -> Self;
#[must_use]
fn dec(self) -> Self;
}
impl<T: AdditionMonoidWithSub + One> IncDec for T {
fn inc(self) -> Self {
self + T::one()
}
fn dec(self) -> Self {
self - T::one()
}
}
impl<T: AdditionMonoidWithSub + One> IncDec for Vec<T> {
fn inc(mut self) -> Self {
self.iter_mut().for_each(|i| *i += T::one());
self
}
fn dec(mut self) -> Self {
self.iter_mut().for_each(|i| *i -= T::one());
self
}
}
impl<T: AdditionMonoidWithSub + One> IncDec for Vec<Vec<T>> {
fn inc(mut self) -> Self {
self.iter_mut().for_each(|v| v.iter_mut().for_each(|i| *i += T::one()));
self
}
fn dec(mut self) -> Self {
self.iter_mut().for_each(|v| v.iter_mut().for_each(|i| *i -= T::one()));
self
}
}
impl<T: AdditionMonoidWithSub + One, U: AdditionMonoidWithSub + One> IncDec
for Vec<(T, U)> {
fn inc(mut self) -> Self {
self.iter_mut()
.for_each(|(i, j)| {
*i += T::one();
*j += U::one();
});
self
}
fn dec(mut self) -> Self {
self.iter_mut()
.for_each(|(i, j)| {
*i -= T::one();
*j -= U::one();
});
self
}
}
impl<T: AdditionMonoidWithSub + One, U: AdditionMonoidWithSub + One, V> IncDec
for Vec<(T, U, V)> {
fn inc(mut self) -> Self {
self.iter_mut()
.for_each(|(i, j, _)| {
*i += T::one();
*j += U::one();
});
self
}
fn dec(mut self) -> Self {
self.iter_mut()
.for_each(|(i, j, _)| {
*i -= T::one();
*j -= U::one();
});
self
}
}
impl<T: AdditionMonoidWithSub + One, U: AdditionMonoidWithSub + One, V, W> IncDec
for Vec<(T, U, V, W)> {
fn inc(mut self) -> Self {
self.iter_mut()
.for_each(|(i, j, ..)| {
*i += T::one();
*j += U::one();
});
self
}
fn dec(mut self) -> Self {
self.iter_mut()
.for_each(|(i, j, ..)| {
*i -= T::one();
*j -= U::one();
});
self
}
}
impl<T: AdditionMonoidWithSub + One, U: AdditionMonoidWithSub + One, V, W, X> IncDec
for Vec<(T, U, V, W, X)> {
fn inc(mut self) -> Self {
self.iter_mut()
.for_each(|(i, j, ..)| {
*i += T::one();
*j += U::one();
});
self
}
fn dec(mut self) -> Self {
self.iter_mut()
.for_each(|(i, j, ..)| {
*i -= T::one();
*j -= U::one();
});
self
}
}
impl<T: AdditionMonoidWithSub + One, U: AdditionMonoidWithSub + One> IncDec for (T, U) {
fn inc(mut self) -> Self {
self.0 += T::one();
self.1 += U::one();
self
}
fn dec(mut self) -> Self {
self.0 -= T::one();
self.1 -= U::one();
self
}
}
}
}
}
pub mod graph {
use crate::algo_lib::collections::dsu::DSU;
use crate::algo_lib::graph::edges::bi_edge::BiEdge;
use crate::algo_lib::graph::edges::edge::Edge;
use crate::algo_lib::graph::edges::edge_trait::{BidirectionalEdgeTrait, EdgeTrait};
use std::ops::{Index, IndexMut};
#[derive(Clone)]
pub struct Graph<E: EdgeTrait> {
edges: Vec<Vec<E>>,
edge_count: usize,
}
impl<E: EdgeTrait> Graph<E> {
pub fn new(vertex_count: usize) -> Self {
Self {
edges: vec![Vec::new(); vertex_count],
edge_count: 0,
}
}
pub fn add_edge(&mut self, (from, mut edge): (usize, E)) -> usize {
let to = edge.to();
assert!(to < self.vertex_count());
let direct_id = self.edges[from].len();
edge.set_id(self.edge_count);
self.edges[from].push(edge);
if E::REVERSABLE {
let rev_id = self.edges[to].len();
self.edges[from][direct_id].set_reverse_id(rev_id);
let mut rev_edge = self.edges[from][direct_id].reverse_edge(from);
rev_edge.set_id(self.edge_count);
rev_edge.set_reverse_id(direct_id);
self.edges[to].push(rev_edge);
}
self.edge_count += 1;
direct_id
}
pub fn add_vertices(&mut self, cnt: usize) {
self.edges.resize(self.edges.len() + cnt, Vec::new());
}
pub fn clear(&mut self) {
self.edge_count = 0;
for ve in self.edges.iter_mut() {
ve.clear();
}
}
pub fn vertex_count(&self) -> usize {
self.edges.len()
}
pub fn edge_count(&self) -> usize {
self.edge_count
}
pub fn degrees(&self) -> Vec<usize> {
self.edges.iter().map(|v| v.len()).collect()
}
}
impl<E: BidirectionalEdgeTrait> Graph<E> {
pub fn is_tree(&self) -> bool {
if self.edge_count + 1 != self.vertex_count() {
false
} else {
self.is_connected()
}
}
pub fn is_forest(&self) -> bool {
let mut dsu = DSU::new(self.vertex_count());
for i in 0..self.vertex_count() {
for e in self[i].iter() {
if i <= e.to() && !dsu.union(i, e.to()) {
return false;
}
}
}
true
}
pub fn is_connected(&self) -> bool {
let mut dsu = DSU::new(self.vertex_count());
for i in 0..self.vertex_count() {
for e in self[i].iter() {
dsu.union(i, e.to());
}
}
dsu.set_count() == 1
}
}
impl<E: EdgeTrait> Index<usize> for Graph<E> {
type Output = [E];
fn index(&self, index: usize) -> &Self::Output {
&self.edges[index]
}
}
impl<E: EdgeTrait> IndexMut<usize> for Graph<E> {
fn index_mut(&mut self, index: usize) -> &mut Self::Output {
&mut self.edges[index]
}
}
impl Graph<Edge<()>> {
pub fn from_edges(n: usize, edges: &[(usize, usize)]) -> Self {
let mut graph = Self::new(n);
for &(from, to) in edges {
graph.add_edge(Edge::new(from, to));
}
graph
}
}
impl<P: Clone> Graph<Edge<P>> {
pub fn from_edges_with_payload(n: usize, edges: &[(usize, usize, P)]) -> Self {
let mut graph = Self::new(n);
for (from, to, p) in edges.iter() {
graph.add_edge(Edge::with_payload(*from, *to, p.clone()));
}
graph
}
}
impl Graph<BiEdge<()>> {
pub fn from_biedges(n: usize, edges: &[(usize, usize)]) -> Self {
let mut graph = Self::new(n);
for &(from, to) in edges {
graph.add_edge(BiEdge::new(from, to));
}
graph
}
}
impl<P: Clone> Graph<BiEdge<P>> {
pub fn from_biedges_with_payload(n: usize, edges: &[(usize, usize, P)]) -> Self {
let mut graph = Self::new(n);
for (from, to, p) in edges.iter() {
graph.add_edge(BiEdge::with_payload(*from, *to, p.clone()));
}
graph
}
}
pub mod edges {
pub mod bi_edge {
use crate::algo_lib::graph::edges::bi_edge_trait::BiEdgeTrait;
use crate::algo_lib::graph::edges::edge_id::{EdgeId, NoId, WithId};
use crate::algo_lib::graph::edges::edge_trait::{BidirectionalEdgeTrait, EdgeTrait};
#[derive(Clone)]
pub struct BiEdgeRaw<Id: EdgeId, P> {
to: u32,
id: Id,
payload: P,
}
impl<Id: EdgeId> BiEdgeRaw<Id, ()> {
pub fn new(from: usize, to: usize) -> (usize, Self) {
(
from,
Self {
to: to as u32,
id: Id::new(),
payload: (),
},
)
}
}
impl<Id: EdgeId, P> BiEdgeRaw<Id, P> {
pub fn with_payload(from: usize, to: usize, payload: P) -> (usize, Self) {
(from, Self::with_payload_impl(to, payload))
}
fn with_payload_impl(to: usize, payload: P) -> BiEdgeRaw<Id, P> {
Self {
to: to as u32,
id: Id::new(),
payload,
}
}
}
impl<Id: EdgeId, P: Clone> BidirectionalEdgeTrait for BiEdgeRaw<Id, P> {}
impl<Id: EdgeId, P: Clone> EdgeTrait for BiEdgeRaw<Id, P> {
type Payload = P;
const REVERSABLE: bool = true;
fn to(&self) -> usize {
self.to as usize
}
fn id(&self) -> usize {
self.id.id()
}
fn set_id(&mut self, id: usize) {
self.id.set_id(id);
}
fn reverse_id(&self) -> usize {
panic!("no reverse id")
}
fn set_reverse_id(&mut self, _: usize) {}
fn reverse_edge(&self, from: usize) -> Self {
Self::with_payload_impl(from, self.payload.clone())
}
fn payload(&self) -> &P {
&self.payload
}
}
impl<Id: EdgeId, P: Clone> BiEdgeTrait for BiEdgeRaw<Id, P> {}
pub type BiEdge<P> = BiEdgeRaw<NoId, P>;
pub type BiEdgeWithId<P> = BiEdgeRaw<WithId, P>;
}
pub mod bi_edge_trait {
use crate::algo_lib::graph::edges::edge_trait::EdgeTrait;
pub trait BiEdgeTrait: EdgeTrait {}
}
pub mod edge {
use crate::algo_lib::graph::edges::edge_id::{EdgeId, NoId, WithId};
use crate::algo_lib::graph::edges::edge_trait::EdgeTrait;
#[derive(Clone)]
pub struct EdgeRaw<Id: EdgeId, P> {
to: u32,
id: Id,
payload: P,
}
impl<Id: EdgeId> EdgeRaw<Id, ()> {
pub fn new(from: usize, to: usize) -> (usize, Self) {
(
from,
Self {
to: to as u32,
id: Id::new(),
payload: (),
},
)
}
}
impl<Id: EdgeId, P> EdgeRaw<Id, P> {
pub fn with_payload(from: usize, to: usize, payload: P) -> (usize, Self) {
(from, Self::with_payload_impl(to, payload))
}
fn with_payload_impl(to: usize, payload: P) -> Self {
Self {
to: to as u32,
id: Id::new(),
payload,
}
}
}
impl<Id: EdgeId, P: Clone> EdgeTrait for EdgeRaw<Id, P> {
type Payload = P;
const REVERSABLE: bool = false;
fn to(&self) -> usize {
self.to as usize
}
fn id(&self) -> usize {
self.id.id()
}
fn set_id(&mut self, id: usize) {
self.id.set_id(id);
}
fn reverse_id(&self) -> usize {
panic!("no reverse")
}
fn set_reverse_id(&mut self, _: usize) {
panic!("no reverse")
}
fn reverse_edge(&self, _: usize) -> Self {
panic!("no reverse")
}
fn payload(&self) -> &P {
&self.payload
}
}
pub type Edge<P> = EdgeRaw<NoId, P>;
pub type EdgeWithId<P> = EdgeRaw<WithId, P>;
}
pub mod edge_id {
pub trait EdgeId: Clone {
fn new() -> Self;
fn id(&self) -> usize;
fn set_id(&mut self, id: usize);
}
#[derive(Clone)]
pub struct WithId {
id: u32,
}
impl EdgeId for WithId {
fn new() -> Self {
Self { id: 0 }
}
fn id(&self) -> usize {
self.id as usize
}
fn set_id(&mut self, id: usize) {
self.id = id as u32;
}
}
#[derive(Clone)]
pub struct NoId {}
impl EdgeId for NoId {
fn new() -> Self {
Self {}
}
fn id(&self) -> usize {
panic!("Id called on no id")
}
fn set_id(&mut self, _: usize) {}
}
}
pub mod edge_trait {
pub trait EdgeTrait: Clone {
type Payload;
const REVERSABLE: bool;
fn to(&self) -> usize;
fn id(&self) -> usize;
fn set_id(&mut self, id: usize);
fn reverse_id(&self) -> usize;
fn set_reverse_id(&mut self, reverse_id: usize);
#[must_use]
fn reverse_edge(&self, from: usize) -> Self;
fn payload(&self) -> &Self::Payload;
}
pub trait BidirectionalEdgeTrait: EdgeTrait {}
}
}
}
pub mod io {
pub mod input {
use crate::algo_lib::collections::vec_ext::default::default_vec;
use std::io::Read;
use std::mem::MaybeUninit;
pub struct Input<'s> {
input: &'s mut (dyn Read + Send),
buf: Vec<u8>,
at: usize,
buf_read: usize,
eol: bool,
}
macro_rules! read_impl {
($t:ty, $read_name:ident, $read_vec_name:ident) => {
pub fn $read_name (& mut self) -> $t { self.read() } pub fn $read_vec_name (& mut
self, len : usize) -> Vec <$t > { self.read_vec(len) }
};
($t:ty, $read_name:ident, $read_vec_name:ident, $read_pair_vec_name:ident) => {
read_impl!($t, $read_name, $read_vec_name); pub fn $read_pair_vec_name (& mut
self, len : usize) -> Vec < ($t, $t) > { self.read_vec(len) }
};
}
impl<'s> Input<'s> {
const DEFAULT_BUF_SIZE: usize = 4096;
pub fn new(input: &'s mut (dyn Read + Send)) -> Self {
Self {
input,
buf: default_vec(Self::DEFAULT_BUF_SIZE),
at: 0,
buf_read: 0,
eol: true,
}
}
pub fn new_with_size(input: &'s mut (dyn Read + Send), buf_size: usize) -> Self {
Self {
input,
buf: default_vec(buf_size),
at: 0,
buf_read: 0,
eol: true,
}
}
pub fn get(&mut self) -> Option<u8> {
if self.refill_buffer() {
let res = self.buf[self.at];
self.at += 1;
if res == b'\r' {
self.eol = true;
if self.refill_buffer() && self.buf[self.at] == b'\n' {
self.at += 1;
}
return Some(b'\n');
}
self.eol = res == b'\n';
Some(res)
} else {
None
}
}
pub fn peek(&mut self) -> Option<u8> {
if self.refill_buffer() {
let res = self.buf[self.at];
Some(if res == b'\r' { b'\n' } else { res })
} else {
None
}
}
pub fn skip_whitespace(&mut self) {
while let Some(b) = self.peek() {
if !b.is_ascii_whitespace() {
return;
}
self.get();
}
}
pub fn next_token(&mut self) -> Option<Vec<u8>> {
self.skip_whitespace();
let mut res = Vec::new();
while let Some(c) = self.get() {
if c.is_ascii_whitespace() {
break;
}
res.push(c);
}
if res.is_empty() { None } else { Some(res) }
}
pub fn is_exhausted(&mut self) -> bool {
self.peek().is_none()
}
pub fn is_empty(&mut self) -> bool {
self.skip_whitespace();
self.is_exhausted()
}
pub fn read<T: Readable>(&mut self) -> T {
T::read(self)
}
pub fn read_vec<T: Readable>(&mut self, size: usize) -> Vec<T> {
let mut res = Vec::with_capacity(size);
for _ in 0..size {
res.push(self.read());
}
res
}
pub fn read_char(&mut self) -> u8 {
self.skip_whitespace();
self.get().unwrap()
}
read_impl!(u32, read_unsigned, read_unsigned_vec);
read_impl!(u64, read_u64, read_u64_vec);
read_impl!(usize, read_size, read_size_vec, read_size_pair_vec);
read_impl!(i32, read_int, read_int_vec, read_int_pair_vec);
read_impl!(i64, read_long, read_long_vec, read_long_pair_vec);
read_impl!(i128, read_i128, read_i128_vec);
fn refill_buffer(&mut self) -> bool {
if self.at == self.buf_read {
self.at = 0;
self.buf_read = self.input.read(&mut self.buf).unwrap();
self.buf_read != 0
} else {
true
}
}
pub fn is_eol(&self) -> bool {
self.eol
}
}
pub trait Readable {
fn read(input: &mut Input) -> Self;
}
impl Readable for u8 {
fn read(input: &mut Input) -> Self {
input.read_char()
}
}
impl<T: Readable> Readable for Vec<T> {
fn read(input: &mut Input) -> Self {
let size = input.read();
input.read_vec(size)
}
}
impl<T: Readable, const SIZE: usize> Readable for [T; SIZE] {
fn read(input: &mut Input) -> Self {
unsafe {
let mut res = MaybeUninit::<[T; SIZE]>::uninit();
for i in 0..SIZE {
let ptr: *mut T = (*res.as_mut_ptr()).as_mut_ptr();
ptr.add(i).write(input.read::<T>());
}
res.assume_init()
}
}
}
macro_rules! read_integer {
($($t:ident)+) => {
$(impl Readable for $t { fn read(input : & mut Input) -> Self { input
.skip_whitespace(); let mut c = input.get().unwrap(); let sgn = match c { b'-' =>
{ c = input.get().unwrap(); true } b'+' => { c = input.get().unwrap(); false } _
=> false, }; let mut res = 0; loop { assert!(c.is_ascii_digit()); res *= 10; let
d = (c - b'0') as $t; if sgn { res -= d; } else { res += d; } match input.get() {
None => break, Some(ch) => { if ch.is_ascii_whitespace() { break; } else { c =
ch; } } } } res } })+
};
}
read_integer!(i8 i16 i32 i64 i128 isize u16 u32 u64 u128 usize);
macro_rules! tuple_readable {
($($name:ident)+) => {
impl <$($name : Readable),+> Readable for ($($name,)+) { fn read(input : & mut
Input) -> Self { ($($name ::read(input),)+) } }
};
}
tuple_readable! {
T
}
tuple_readable! {
T U
}
tuple_readable! {
T U V
}
tuple_readable! {
T U V X
}
tuple_readable! {
T U V X Y
}
tuple_readable! {
T U V X Y Z
}
tuple_readable! {
T U V X Y Z A
}
tuple_readable! {
T U V X Y Z A B
}
tuple_readable! {
T U V X Y Z A B C
}
tuple_readable! {
T U V X Y Z A B C D
}
tuple_readable! {
T U V X Y Z A B C D E
}
tuple_readable! {
T U V X Y Z A B C D E F
}
impl Read for Input<'_> {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
if self.at == self.buf_read {
self.input.read(buf)
} else {
let mut i = 0;
while i < buf.len() && self.at < self.buf_read {
buf[i] = self.buf[self.at];
i += 1;
self.at += 1;
}
Ok(i)
}
}
}
}
pub mod output {
use crate::algo_lib::collections::vec_ext::default::default_vec;
use std::cmp::Reverse;
use std::io::{stderr, Stderr, Write};
#[derive(Copy, Clone)]
pub enum BoolOutput {
YesNo,
YesNoCaps,
PossibleImpossible,
Custom(&'static str, &'static str),
}
impl BoolOutput {
pub fn output(&self, output: &mut Output, val: bool) {
(if val { self.yes() } else { self.no() }).write(output);
}
fn yes(&self) -> &str {
match self {
BoolOutput::YesNo => "Yes",
BoolOutput::YesNoCaps => "YES",
BoolOutput::PossibleImpossible => "Possible",
BoolOutput::Custom(yes, _) => yes,
}
}
fn no(&self) -> &str {
match self {
BoolOutput::YesNo => "No",
BoolOutput::YesNoCaps => "NO",
BoolOutput::PossibleImpossible => "Impossible",
BoolOutput::Custom(_, no) => no,
}
}
}
pub struct Output<'s> {
output: &'s mut dyn Write,
buf: Vec<u8>,
at: usize,
auto_flush: bool,
bool_output: BoolOutput,
precision: Option<usize>,
separator: u8,
}
impl<'s> Output<'s> {
const DEFAULT_BUF_SIZE: usize = 4096;
pub fn new(output: &'s mut dyn Write) -> Self {
Self {
output,
buf: default_vec(Self::DEFAULT_BUF_SIZE),
at: 0,
auto_flush: false,
bool_output: BoolOutput::YesNoCaps,
precision: None,
separator: b' ',
}
}
pub fn new_with_auto_flush(output: &'s mut dyn Write) -> Self {
Self {
output,
buf: default_vec(Self::DEFAULT_BUF_SIZE),
at: 0,
auto_flush: true,
bool_output: BoolOutput::YesNoCaps,
precision: None,
separator: b' ',
}
}
pub fn flush(&mut self) {
if self.at != 0 {
self.output.write_all(&self.buf[..self.at]).unwrap();
self.output.flush().unwrap();
self.at = 0;
}
}
pub fn print<T: Writable>(&mut self, s: T) {
s.write(self);
self.maybe_flush();
}
pub fn print_line<T: Writable>(&mut self, s: T) {
self.print(s);
self.put(b'\n');
self.maybe_flush();
}
pub fn put(&mut self, b: u8) {
self.buf[self.at] = b;
self.at += 1;
if self.at == self.buf.len() {
self.flush();
}
}
pub fn maybe_flush(&mut self) {
if self.auto_flush {
self.flush();
}
}
pub fn print_per_line<T: Writable>(&mut self, arg: &[T]) {
self.print_per_line_iter(arg.iter());
}
pub fn print_iter<T: Writable, I: Iterator<Item = T>>(&mut self, iter: I) {
let mut first = true;
for e in iter {
if first {
first = false;
} else {
self.put(self.separator);
}
e.write(self);
}
}
pub fn print_line_iter<T: Writable, I: Iterator<Item = T>>(&mut self, iter: I) {
self.print_iter(iter);
self.put(b'\n');
}
pub fn print_per_line_iter<T: Writable, I: Iterator<Item = T>>(&mut self, iter: I) {
for e in iter {
e.write(self);
self.put(b'\n');
}
}
pub fn set_bool_output(&mut self, bool_output: BoolOutput) {
self.bool_output = bool_output;
}
pub fn set_precision(&mut self, precision: usize) {
self.precision = Some(precision);
}
pub fn reset_precision(&mut self) {
self.precision = None;
}
pub fn get_precision(&self) -> Option<usize> {
self.precision
}
pub fn separator(&self) -> u8 {
self.separator
}
pub fn set_separator(&mut self, separator: u8) {
self.separator = separator;
}
}
impl Write for Output<'_> {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
let mut start = 0usize;
let mut rem = buf.len();
while rem > 0 {
let len = (self.buf.len() - self.at).min(rem);
self.buf[self.at..self.at + len].copy_from_slice(&buf[start..start + len]);
self.at += len;
if self.at == self.buf.len() {
self.flush();
}
start += len;
rem -= len;
}
self.maybe_flush();
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
self.flush();
Ok(())
}
}
pub trait Writable {
fn write(&self, output: &mut Output);
}
impl Writable for &str {
fn write(&self, output: &mut Output) {
output.write_all(self.as_bytes()).unwrap();
}
}
impl Writable for String {
fn write(&self, output: &mut Output) {
output.write_all(self.as_bytes()).unwrap();
}
}
impl Writable for char {
fn write(&self, output: &mut Output) {
output.put(*self as u8);
}
}
impl Writable for u8 {
fn write(&self, output: &mut Output) {
output.put(*self);
}
}
impl<T: Writable> Writable for [T] {
fn write(&self, output: &mut Output) {
output.print_iter(self.iter());
}
}
impl<T: Writable, const N: usize> Writable for [T; N] {
fn write(&self, output: &mut Output) {
output.print_iter(self.iter());
}
}
impl<T: Writable + ?Sized> Writable for &T {
fn write(&self, output: &mut Output) {
T::write(self, output)
}
}
impl<T: Writable> Writable for Vec<T> {
fn write(&self, output: &mut Output) {
self.as_slice().write(output);
}
}
impl Writable for () {
fn write(&self, _output: &mut Output) {}
}
macro_rules! write_to_string {
($($t:ident)+) => {
$(impl Writable for $t { fn write(& self, output : & mut Output) { self
.to_string().write(output); } })+
};
}
write_to_string!(u16 u32 u64 u128 usize i8 i16 i32 i64 i128 isize);
macro_rules! tuple_writable {
($name0:ident $($name:ident : $id:tt)*) => {
impl <$name0 : Writable, $($name : Writable,)*> Writable for ($name0, $($name,)*)
{ fn write(& self, out : & mut Output) { self.0.write(out); $(out.put(out
.separator); self.$id .write(out);)* } }
};
}
tuple_writable! {
T
}
tuple_writable! {
T U : 1
}
tuple_writable! {
T U : 1 V : 2
}
tuple_writable! {
T U : 1 V : 2 X : 3
}
tuple_writable! {
T U : 1 V : 2 X : 3 Y : 4
}
tuple_writable! {
T U : 1 V : 2 X : 3 Y : 4 Z : 5
}
tuple_writable! {
T U : 1 V : 2 X : 3 Y : 4 Z : 5 A : 6
}
tuple_writable! {
T U : 1 V : 2 X : 3 Y : 4 Z : 5 A : 6 B : 7
}
tuple_writable! {
T U : 1 V : 2 X : 3 Y : 4 Z : 5 A : 6 B : 7 C : 8
}
impl<T: Writable> Writable for Option<T> {
fn write(&self, output: &mut Output) {
match self {
None => (-1).write(output),
Some(t) => t.write(output),
}
}
}
impl Writable for bool {
fn write(&self, output: &mut Output) {
let bool_output = output.bool_output;
bool_output.output(output, *self)
}
}
impl<T: Writable> Writable for Reverse<T> {
fn write(&self, output: &mut Output) {
self.0.write(output);
}
}
static mut ERR: Option<Stderr> = None;
pub fn err() -> Output<'static> {
unsafe {
if ERR.is_none() {
ERR = Some(stderr());
}
Output::new_with_auto_flush(ERR.as_mut().unwrap())
}
}
}
}
pub mod misc {
pub mod direction {
#[derive(Copy, Clone)]
pub enum Direction {
Left,
Right,
}
}
pub mod extensions {
pub mod replace_with {
use std::{mem, ptr};
pub trait ReplaceWith<F: FnOnce(Self) -> Self>: Sized {
fn replace_with(&mut self, f: F) {
unsafe {
let old = ptr::read(self);
let new = on_unwind(move || f(old), || std::process::abort());
ptr::write(self, new);
}
}
}
impl<T, F: FnOnce(Self) -> Self> ReplaceWith<F> for T {}
pub fn on_unwind<F: FnOnce() -> T, T, P: FnOnce()>(f: F, p: P) -> T {
let x = OnDrop(mem::ManuallyDrop::new(p));
let t = f();
let mut x = mem::ManuallyDrop::new(x);
unsafe { mem::ManuallyDrop::drop(&mut x.0) };
t
}
struct OnDrop<F: FnOnce()>(mem::ManuallyDrop<F>);
impl<F: FnOnce()> Drop for OnDrop<F> {
#[inline(always)]
fn drop(&mut self) {
(unsafe { ptr::read(&*self.0) })();
}
}
}
}
pub mod random {
use crate::algo_lib::collections::slice_ext::indices::Indices;
use crate::algo_lib::numbers::num_traits::algebra::IntegerSemiRingWithSub;
use crate::algo_lib::numbers::num_traits::ord::MinMax;
use crate::algo_lib::numbers::num_traits::primitive::Primitive;
use std::cell::RefCell;
use std::ops::{RangeBounds, Rem};
use std::time::SystemTime;
const NN: usize = 312;
const MM: usize = 156;
const MATRIX_A: u64 = 0xB5026F5AA96619E9;
const UM: u64 = 0xFFFFFFFF80000000;
const LM: u64 = 0x7FFFFFFF;
const F: u64 = 6364136223846793005;
const MAG01: [u64; 2] = [0, MATRIX_A];
pub struct Random {
mt: [u64; NN],
index: usize,
}
impl Random {
pub fn new() -> Self {
Self::new_with_seed(
(SystemTime::UNIX_EPOCH.elapsed().unwrap().as_nanos() & 0xFFFFFFFFFFFFFFFF)
as u64,
)
}
pub fn new_with_seed(seed: u64) -> Self {
let mut res = Self { mt: [0u64; NN], index: NN };
res.mt[0] = seed;
for i in 1..NN {
res.mt[i] = F
.wrapping_mul(res.mt[i - 1] ^ (res.mt[i - 1] >> 62))
.wrapping_add(i as u64);
}
res
}
fn gen_impl(&mut self) -> u64 {
if self.index == NN {
for i in 0..(NN - MM) {
let x = (self.mt[i] & UM) | (self.mt[i + 1] & LM);
self.mt[i] = self.mt[i + MM] ^ (x >> 1) ^ MAG01[(x & 1) as usize];
}
for i in (NN - MM)..(NN - 1) {
let x = (self.mt[i] & UM) | (self.mt[i + 1] & LM);
self.mt[i] = self.mt[i + MM - NN] ^ (x >> 1) ^ MAG01[(x & 1) as usize];
}
let x = (self.mt[NN - 1] & UM) | (self.mt[0] & LM);
self.mt[NN - 1] = self.mt[MM - 1] ^ (x >> 1) ^ MAG01[(x & 1) as usize];
self.index = 0;
}
let mut x = self.mt[self.index];
self.index += 1;
x ^= (x >> 29) & 0x5555555555555555;
x ^= (x << 17) & 0x71D67FFFEDA60000;
x ^= (x << 37) & 0xFFF7EEE000000000;
x ^= x >> 43;
x
}
pub fn gen<T>(&mut self) -> T
where
u64: Primitive<T>,
{
self.gen_impl().to()
}
pub fn gen_u128(&mut self) -> u128 {
(self.gen_impl() as u128) << 64 | self.gen_impl() as u128
}
pub fn gen_i128(&mut self) -> i128 {
self.gen_u128() as i128
}
pub fn gen_bool(&mut self) -> bool {
(self.gen_impl() & 1) == 1
}
pub fn gen_bound<T: Rem<Output = T> + Primitive<u64>>(&mut self, n: T) -> T
where
u64: Primitive<T>,
{
(self.gen_impl() % n.to()).to()
}
pub fn gen_range<T: IntegerSemiRingWithSub + Primitive<u64> + MinMax>(
&mut self,
range: impl RangeBounds<T>,
) -> T
where
u64: Primitive<T>,
{
let f = match range.start_bound() {
std::ops::Bound::Included(&s) => s,
std::ops::Bound::Excluded(&s) => s + T::one(),
std::ops::Bound::Unbounded => T::min_val(),
};
let t = match range.end_bound() {
std::ops::Bound::Included(&e) => e,
std::ops::Bound::Excluded(&e) => e - T::one(),
std::ops::Bound::Unbounded => T::max_val(),
};
if f == T::min_val() && t == T::max_val() {
self.gen()
} else {
f + self.gen_bound(t - f + T::one())
}
}
}
thread_local! {
static RANDOM : RefCell < Random > = RefCell::new(Random::new());
}
pub fn gen<T>() -> T
where
u64: Primitive<T>,
{
RANDOM.with_borrow_mut(|r| r.gen())
}
pub fn gen_u128() -> u128 {
RANDOM.with_borrow_mut(|r| r.gen_u128())
}
pub fn gen_i128() -> i128 {
RANDOM.with_borrow_mut(|r| r.gen_i128())
}
pub fn gen_bool() -> bool {
RANDOM.with_borrow_mut(|r| r.gen_bool())
}
pub fn gen_bound<T: Rem<Output = T> + Primitive<u64>>(n: T) -> T
where
u64: Primitive<T>,
{
RANDOM.with_borrow_mut(|r| r.gen_bound(n))
}
pub fn gen_range<T: IntegerSemiRingWithSub + Primitive<u64> + MinMax>(
range: impl RangeBounds<T>,
) -> T
where
u64: Primitive<T>,
{
RANDOM.with_borrow_mut(|r| r.gen_range(range))
}
pub trait Shuffle {
fn shuffle(&mut self);
}
impl<T> Shuffle for [T] {
fn shuffle(&mut self) {
for i in self.indices() {
let at = gen_bound(i + 1);
self.swap(i, at);
}
}
}
}
pub mod recursive_function {
use std::marker::PhantomData;
macro_rules! recursive_function {
($name:ident, $trait:ident, ($($type:ident $arg:ident,)*)) => {
pub trait $trait <$($type,)* Output > { fn call(& mut self, $($arg : $type,)*) ->
Output; } pub struct $name < F, $($type,)* Output > where F : FnMut(& mut dyn
$trait <$($type,)* Output >, $($type,)*) -> Output, { f : std::cell::UnsafeCell <
F >, $($arg : PhantomData <$type >,)* phantom_output : PhantomData < Output >, }
impl < F, $($type,)* Output > $name < F, $($type,)* Output > where F : FnMut(&
mut dyn $trait <$($type,)* Output >, $($type,)*) -> Output, { pub fn new(f : F)
-> Self { Self { f : std::cell::UnsafeCell::new(f), $($arg :
Default::default(),)* phantom_output : Default::default(), } } } impl < F,
$($type,)* Output > $trait <$($type,)* Output > for $name < F, $($type,)* Output
> where F : FnMut(& mut dyn $trait <$($type,)* Output >, $($type,)*) -> Output, {
fn call(& mut self, $($arg : $type,)*) -> Output { unsafe { (* self.f.get())
(self, $($arg,)*) } } }
};
}
recursive_function!(RecursiveFunction0, Callable0, ());
recursive_function!(RecursiveFunction, Callable, (Arg arg,));
recursive_function!(RecursiveFunction2, Callable2, (Arg1 arg1, Arg2 arg2,));
recursive_function!(RecursiveFunction3, Callable3, (Arg1 arg1, Arg2 arg2, Arg3 arg3,));
recursive_function!(
RecursiveFunction4, Callable4, (Arg1 arg1, Arg2 arg2, Arg3 arg3, Arg4 arg4,)
);
recursive_function!(
RecursiveFunction5, Callable5, (Arg1 arg1, Arg2 arg2, Arg3 arg3, Arg4 arg4, Arg5
arg5,)
);
recursive_function!(
RecursiveFunction6, Callable6, (Arg1 arg1, Arg2 arg2, Arg3 arg3, Arg4 arg4, Arg5
arg5, Arg6 arg6,)
);
recursive_function!(
RecursiveFunction7, Callable7, (Arg1 arg1, Arg2 arg2, Arg3 arg3, Arg4 arg4, Arg5
arg5, Arg6 arg6, Arg7 arg7,)
);
recursive_function!(
RecursiveFunction8, Callable8, (Arg1 arg1, Arg2 arg2, Arg3 arg3, Arg4 arg4, Arg5
arg5, Arg6 arg6, Arg7 arg7, Arg8 arg8,)
);
recursive_function!(
RecursiveFunction9, Callable9, (Arg1 arg1, Arg2 arg2, Arg3 arg3, Arg4 arg4, Arg5
arg5, Arg6 arg6, Arg7 arg7, Arg8 arg8, Arg9 arg9,)
);
}
pub mod test_type {
pub enum TestType {
Single,
MultiNumber,
MultiEof,
}
pub enum TaskType {
Classic,
Interactive,
}
}
}
pub mod numbers {
pub mod num_traits {
pub mod algebra {
use crate::algo_lib::numbers::num_traits::invertible::Invertible;
use std::ops::{
Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Rem, RemAssign, Sub, SubAssign,
};
pub trait Zero {
fn zero() -> Self;
}
pub trait One {
fn one() -> Self;
}
pub trait AdditionMonoid: Add<Output = Self> + AddAssign + Zero + Eq + Sized {}
impl<T: Add<Output = Self> + AddAssign + Zero + Eq> AdditionMonoid for T {}
pub trait AdditionMonoidWithSub: AdditionMonoid + Sub<Output = Self> + SubAssign {}
impl<T: AdditionMonoid + Sub<Output = Self> + SubAssign> AdditionMonoidWithSub for T {}
pub trait AdditionGroup: AdditionMonoidWithSub + Neg<Output = Self> {}
impl<T: AdditionMonoidWithSub + Neg<Output = Self>> AdditionGroup for T {}
pub trait MultiplicationMonoid: Mul<Output = Self> + MulAssign + One + Eq + Sized {}
impl<T: Mul<Output = Self> + MulAssign + One + Eq> MultiplicationMonoid for T {}
pub trait IntegerMultiplicationMonoid: MultiplicationMonoid + Div<
Output = Self,
> + Rem<Output = Self> + DivAssign + RemAssign {}
impl<
T: MultiplicationMonoid + Div<Output = Self> + Rem<Output = Self> + DivAssign
+ RemAssign,
> IntegerMultiplicationMonoid for T {}
pub trait MultiplicationGroup: MultiplicationMonoid + Div<
Output = Self,
> + DivAssign + Invertible<Output = Self> {}
impl<
T: MultiplicationMonoid + Div<Output = Self> + DivAssign + Invertible<Output = Self>,
> MultiplicationGroup for T {}
pub trait SemiRing: AdditionMonoid + MultiplicationMonoid {}
impl<T: AdditionMonoid + MultiplicationMonoid> SemiRing for T {}
pub trait SemiRingWithSub: AdditionMonoidWithSub + SemiRing {}
impl<T: AdditionMonoidWithSub + SemiRing> SemiRingWithSub for T {}
pub trait Ring: SemiRing + AdditionGroup {}
impl<T: SemiRing + AdditionGroup> Ring for T {}
pub trait IntegerSemiRing: SemiRing + IntegerMultiplicationMonoid {}
impl<T: SemiRing + IntegerMultiplicationMonoid> IntegerSemiRing for T {}
pub trait IntegerSemiRingWithSub: SemiRingWithSub + IntegerSemiRing {}
impl<T: SemiRingWithSub + IntegerSemiRing> IntegerSemiRingWithSub for T {}
pub trait IntegerRing: IntegerSemiRing + Ring {}
impl<T: IntegerSemiRing + Ring> IntegerRing for T {}
pub trait Field: Ring + MultiplicationGroup {}
impl<T: Ring + MultiplicationGroup> Field for T {}
macro_rules! zero_one_integer_impl {
($($t:ident)+) => {
$(impl Zero for $t { fn zero() -> Self { 0 } } impl One for $t { fn one() -> Self
{ 1 } })+
};
}
zero_one_integer_impl!(i128 i64 i32 i16 i8 isize u128 u64 u32 u16 u8 usize);
}
pub mod invertible {
pub trait Invertible {
type Output;
fn inv(&self) -> Option<Self::Output>;
}
}
pub mod ord {
pub trait MinMax: PartialOrd {
fn min_val() -> Self;
fn max_val() -> Self;
}
macro_rules! min_max_integer_impl {
($($t:ident)+) => {
$(impl MinMax for $t { fn min_val() -> Self { std::$t ::MIN } fn max_val() ->
Self { std::$t ::MAX } })+
};
}
min_max_integer_impl!(i128 i64 i32 i16 i8 isize u128 u64 u32 u16 u8 usize);
}
pub mod primitive {
pub trait Primitive<T>: Copy {
fn to(self) -> T;
fn from(t: T) -> Self;
}
macro_rules! primitive_one {
($t:ident, $($u:ident)+) => {
$(impl Primitive <$u > for $t { fn to(self) -> $u { self as $u } fn from(t : $u)
-> Self { t as $t } })+
};
}
macro_rules! primitive {
($($t:ident)+) => {
$(primitive_one!($t, u8 u16 u32 u64 u128 usize i8 i16 i32 i64 i128 isize);)+
};
}
primitive!(u8 u16 u32 u64 u128 usize i8 i16 i32 i64 i128 isize);
}
}
}
}