AOJ 2829 Room Assignment

問題

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2829

解法

JAG の wiki に分かりやすい解説があります。

2017/Practice/模擬国内予選/講評 - ACM-ICPC Japanese Alumni Group

場合分けとしては以下の 3 つです

• 長さ 3 以上の閉路が存在する場合は 0
• 閉路を潰した時に、全てのツリーの葉の数が 2 以下でなければ 0
• 長さ 3 以上の弱連結成分が存在しない時
• それ以外

コード

import java.util.Scanner

import scala.collection.mutable.ArrayBuffer

object Main {
  private val Mod = (1e9 + 7).toInt

  // 階乗を前計算しておく
  private val Fact = {
    var cur = 1L
    for (i <- 0 to 200000) yield {
      if (i == 0) {
        1
      } else {
        cur = (cur * i) % Mod
        cur
      }
    }
  }

  // 2 の累乗を前計算しておく
  private val Pow = {
    val pow = new Array[Int](200000)
    pow(0) = 1
    for (i <- 1 until pow.length) pow(i) = pow(i - 1) * 2 % Mod
    pow
  }

  def main(args: Array[String]): Unit = {
    val in = new Scanner(System.in)
    while (true) {
      val n = in.nextInt()
      if (n == 0) {
        return
      }
      val a = for (_ <- 0 until n) yield {
        in.nextInt() - 1
      }
      println(solve(a.toArray))
    }
  }

  // 長さ 3 以上の閉路が存在しないことを確認する
  def cycleCheck(graph: IndexedSeq[Seq[Int]]): Boolean = {
    val cmp = StronglyConnectedComponents.decompose(graph)
    var map = Map[Int, Int]()
    cmp.foreach { c => map += (c -> (map.getOrElse(c, 0) + 1)) }
    map.values.forall(count => count <= 2)
  }

  // 全ての弱連結成分がパスが閉路を潰すとパスになることを確認する
  def pathCheck(graph: IndexedSeq[Seq[Int]]): Boolean = {
    val uf = new UnionFind(graph.size)
    for {
      i <- graph.indices
      j <- graph(i)
    } uf.unite(i, j)

    var map = Map[Int, Int]()
    for {
      i <- graph.indices
      if graph(i).isEmpty
    } {
      val count = map.getOrElse(uf.find(i), 0)
      map += (uf.find(i) -> (count + 1))
    }
    map.values.forall(count => count <= 2)
  }

  // 要素数 2 の弱連結成分の数と、要素数 3 以上の弱連結成分の数を算出する
  def getSize(array: Array[Int]): (Int, Int) = {
    val N = array.length
    val uf = new UnionFind(N)
    for (i <- array.indices) uf.unite(i, array(i))
    val sizes = for {
      i <- 0 until N
      if i == 0 || !uf.isSame(i, 0)
    } yield {
      if (i == 0) {
        uf.partialSizeOf(i)
      } else {
        val s = uf.partialSizeOf(i)
        uf.unite(i, 0)
        s
      }
    }

    var three = 0
    var two = 0
    sizes.foreach(size => if (size == 2) {
      two += 1
    } else {
      three += 1
    })

    (two, three)
  }

  def solve(parent: Array[Int]): Int = {
    val n = parent.length
    val combination = new Combination(n + 1, Mod)

    val graph = parent.indices.map(i => new ArrayBuffer[Int]())
    parent.indices.foreach(i => graph(parent(i)).append(i))

    if (cycleCheck(graph) && pathCheck(graph)) {
      val (two, three) = getSize(parent)
      if (three == 0) {
        // 長さ 3 以上の弱連結成分が存在しないとき
        var cur: Long = ((two + 2) / 2) % Mod
        cur = cur * Pow(two) % Mod
        cur = cur * Fact(two) % Mod
        cur.toInt
      } else {
        var ans = 0L
        for (i <- 0 to math.min(two + three - 2, two)) {
          var cur = combination.get(two + three - i - 2, two - i).toLong
          cur = cur * Fact(two) % Mod
          cur = cur * Fact(three) % Mod
          cur = cur * Pow(two + three) % Mod
          cur = cur * (two + three - i - 1) % Mod
          cur = cur * ((i + 2) / 2) % Mod
          ans = (ans + cur) % Mod
        }
        for (i <- 0 to math.min(two + three - 1, two)) {
          var cur = combination.get(two + three - i - 1, two - i).toLong
          cur = cur * Fact(two) % Mod
          cur = cur * Fact(three) % Mod
          cur = cur * Pow(two + three) % Mod
          cur = cur * ((i + 2) / 2) % Mod
          ans = (ans + cur * 2) % Mod
        }
        ans.toInt
      }
    } else {
      0
    }
  }

  class UnionFind(n: Int) {
    private val parent = (0 until n).toArray
    private val sizes = Array.fill[Int](n)(1)
    private var _size = n

    def find(x: Int): Int = {
      if (x == parent(x)) {
        x
      } else {
        parent(x) = find(parent(x))
        parent(x)
      }
    }

    def unite(a: Int, b: Int): Boolean = {
      val fa = find(a)
      val fb = find(b)
      if (fa == fb) {
        false
      } else {
        val (x, y) = if (sizes(fa) >= sizes(fb)) {
          (fa, fb)
        } else {
          (fb, fa)
        }

        parent(y) = x
        sizes(x) += sizes(y)
        sizes(y) = 0

        _size -= 1
        true
      }
    }

    def isSame(x: Int, y: Int): Boolean = find(x) == find(y)

    def partialSizeOf(x: Int): Int = sizes(find(x))

    def size(): Int = _size
  }

  class Combination(max: Int, mod: Int) {
    private val inv = new Array[Long](max + 1)
    private val fact = new Array[Long](max + 1)
    private val invFact = new Array[Long](max + 1)
    inv(1) = 1
    for (i <- 2 to max) inv(i) = inv(mod % i) * (mod - mod / i) % mod
    fact(0) = 1
    invFact(0) = 1
    for (i <- 1 to max) fact(i) = (fact(i - 1) * i) % mod
    for (i <- 1 to max) invFact(i) = (invFact(i - 1) * inv(i)) % mod

    /**
      * get nCm
      */
    def get(n: Int, m: Int): Int = {
      if (n < m) {
        0
      } else fact(n) * invFact(m) % mod * invFact(n - m) % mod
    }.toInt
  }

  object StronglyConnectedComponents {
    def decompose(graph: IndexedSeq[Seq[Int]]): Array[Int] = {
      val vs = new ArrayBuffer[Int]()
      val V = graph.size
      val cmp = new Array[Int](V)

      val rg = graph.indices.map(_ => new ArrayBuffer[Int]())
      for {
        from <- graph.indices
        to <- graph(from)
      } rg(to).append(from)

      var used = Array.fill[Boolean](V)(false)
      var stack = List[Int]()
      val added = Array.fill[Boolean](V)(false)
      for {
        i <- used.indices
        if !used(i)
      } {
        stack = i :: stack
        while (stack.nonEmpty) {
          val v = stack.head
          used(v) = true
          var pushed = false
          for {
            u <- graph(v).reverse
            if !used(u)
          } {
            stack = u :: stack
            pushed = true
          }
          if (!pushed) {
            stack = stack.tail
            if (!added(v)) {
              vs.append(v)
              added(v) = true
            }
          }
        }
      }

      used = Array.fill[Boolean](V)(false)
      var k = 0
      for {
        i <- vs.reverse
        if !used(i)
      } {
        stack = i :: stack
        used(i) = true
        cmp(i) = k
        while (stack.nonEmpty) {
          val v = stack.head
          var pushed = false
          for {
            u <- rg(v)
            if !used(u)
          } {
            used(u) = true
            cmp(u) = k
            stack = u :: stack
            pushed = true
          }
          if (!pushed) {
            stack = stack.tail
          }
        }
        k += 1
      }
      cmp
    }
  }

}

AOJ 3019 Picnic

問題

Picnic | Aizu Online Judge

解法

ワーシャルフロイド→巡回セールスマン→半分全列挙→個数制約付きナップサック

コード

import java.util

import scala.io.StdIn
import scala.util.Try

object Main extends App {
  val INF: Long = 1e15.toLong
  val (n, x, y) = {
    val in = StdIn.readLine().split(" ").map(_.toInt)
    (in(0), in(1), in(2))
  }
  val sweets = for (_ <- 0 until n) yield {
    val k = StdIn.readLine().toInt
    for (_ <- 0 until k) yield {
      val in = StdIn.readLine().split(" ").map(_.toInt)
      Sweet(in(0), in(1), in(2))
    }
  }

  val dist = for (_ <- 0 until n) yield StdIn.readLine().split(" ").map(_.toLong)
  for (k <- 0 until n) for (i <- 0 until n) for (j <- 0 until n) dist(i)(j) = math.min(dist(i)(j), dist(i)(k) + dist(k)(j))

  val cost = Array.fill[Long](1 << n, n)(INF)
  cost(0)(0) = 0
  for (mask <- 0 until 1 << n)
    for (cur <- 0 until n)
      for (next <- 0 until n)
        if ((mask & (1 << next)) == 0) {
          val add = dist(cur)(next)
          cost(mask | (1 << next))(next) = math.min(cost(mask | (1 << next))(next), cost(mask)(cur) + add)
        }


  def knapsack(dp: Array[Long], available: Seq[Sweet]): Unit = {
    for {
      sweet <- available
      money <- 0 until sweet.price
    } {
      val deque = new util.ArrayDeque[(Int, Long)]()
      var num = 0
      while (num * sweet.price + money <= y) {
        val next = num * sweet.price + money


        val v = dp(next) - num * sweet.value
        while (!deque.isEmpty && deque.peekLast()._2 <= v) deque.pollLast()
        deque.addLast((num, v))

        dp(next) = deque.peekFirst()._2 + num * sweet.value
        if (deque.peekFirst()._1 == num - sweet.remain) deque.pollFirst()

        num += 1
      }
    }
    for (i <- 1 until dp.length) dp(i) = math.max(dp(i), dp(i - 1))
  }


  val prefix = n / 2
  val suffix = n - prefix

  val prefixDp = Array.fill[Long](1 << prefix, y + 1)(0)
  for (prefixMask <- 0 until 1 << prefix) {
    val mask = prefixMask << suffix
    val dp = prefixDp(prefixMask)
    val available = for {
      i <- 0 until n
      if (mask & (1 << i)) != 0
      sweet <- sweets(i)
    } yield sweet
    knapsack(dp, available)
  }

  val suffixDp = Array.fill[Long](1 << suffix, y + 1)(0)
  for (mask <- 0 until 1 << suffix) {
    val dp = suffixDp(mask)
    val available = for {
      i <- 0 until n
      if (mask & (1 << i)) != 0
      sweet <- sweets(i)
    } yield sweet
    knapsack(dp, available)
  }
  var ans = 0L
  for (mask <- 0 until 1 << n) {
    val suffixMask = mask & ((1 << suffix) - 1)
    val prefixMask = mask >> suffix
    val travelingCost = (for (last <- 0 until n) yield {
      cost(mask)(last) + dist(last)(0)
    }).min
    val money = math.min(x - travelingCost, y).toInt
    for (prefixMoney <- 0 to money) {
      val suffixMoney = money - prefixMoney
      val value = prefixDp(prefixMask)(prefixMoney) + suffixDp(suffixMask)(suffixMoney)
      ans = math.max(ans, value)
    }
  }
  println(ans)
}

case class Sweet(price: Int, value: Int, remain: Int)

AIM Tech Round 4 (Div. 1) D. Dynamic Shortest Path

問題

http://codeforces.com/contest/843/problem/D

解法

最初にダイクストラで距離を求めておき、各クエリごとにダイクストラで追加で増えた分の距離を計算して加算していく。各クエリごとに距離は 1 しか増えないので、キューを距離の数だけ用意すれば、クエリごとのダイクストラの計算量は になる。

よって、最初のダイクストラが クエリごとのダイクストラが全部で となるので、全体で で とかなので、 10 秒あれば間に合う（いいえ）。

コード

use std::io;
use std::str;
use std::usize;
use std::collections::BinaryHeap;
use std::cmp::Ordering;
use std::collections::BTreeMap;
use std::collections::VecDeque;
use std::cmp;

static INF: i64 = 1 << 60;
static MAX_REQ: usize = 2000000;

fn main() {
    let start = 0;
    let (n, m, q) = {
        let v = read_values::<usize>();
        (v[0], v[1], v[2])
    };

    let mut graph: Vec<Vec<Edge>> = vec![Vec::new(); n];
    let mut edges = Vec::new();
    for _ in 0..m {
        let (from, to, cost) = {
            let v = read_values::<usize>();
            (v[0] - 1, v[1] - 1, v[2] as i64)
        };
        edges.push((from, graph[from].len()));
        graph[from].push(Edge { to: to, cost: cost });
    }

    let mut shortest_dist = dijkstra(start, &graph);
    let mut deque = vec![VecDeque::new(); MAX_REQ];
    let mut add = vec![INF; n];

    let mut modify_queries = Vec::new();

    for _ in 0..q {
        let queries = read_values::<usize>();
        if queries[0] == 1 {
            let to = queries[1] - 1;
            let num = modify_queries.len();

            for edge_idx in &modify_queries {
                let (from, idx): (usize, usize) = edges[*edge_idx];
                unsafe { graph[from].get_unchecked_mut(idx).cost += 1; }
            }
            modify_queries.clear();

            add[start] = 0;
            deque[0].push_back(start);

            for dist in 0..(num + 1) {
                while !deque[dist].is_empty() {
                    let cur = deque[dist].pop_front().unwrap();
                    if add[cur] != dist as i64 {
                        continue;
                    }
                    for e in &graph[cur] {
                        let d = shortest_dist[cur] + add[cur] - shortest_dist[e.to] + e.cost;
                        if d as usize <= num && add[e.to] > d {
                            add[e.to] = d;
                            deque[d as usize].push_back(e.to);
                        }
                    }
                }
            }

            for i in 0..n {
                shortest_dist[i] = cmp::min(shortest_dist[i] + add[i], INF);
                add[i] = INF;
            }

            if shortest_dist[to] >= INF {
                println!("-1");
            } else {
                println!("{}", shortest_dist[to]);
            }
        } else {
            let num = queries[1];
            for i in 0..num { modify_queries.push(queries[i + 2] - 1); }
        }
    }
}

fn dijkstra(start: usize, graph: &Vec<Vec<Edge>>) -> Vec<i64> {
    let n = graph.len();
    let mut heap = BinaryHeap::<Edge>::new();
    heap.push(Edge { to: start, cost: 0 });
    let mut dist = vec![INF; n];
    dist[start] = 0;
    while !heap.is_empty() {
        let p: Edge = heap.pop().unwrap();
        for e in &graph[p.to] {
            if dist[e.to] > e.cost + dist[p.to] {
                dist[e.to] = e.cost + dist[p.to];
                heap.push(Edge { to: e.to, cost: dist[e.to] });
            }
        }
    }
    dist
}

#[derive(Copy, Clone, Eq, PartialEq)]
struct Edge {
    to: usize,
    cost: i64,
}

impl Ord for Edge {
    fn cmp(&self, other: &Edge) -> Ordering {
        other.cost.cmp(&self.cost)
    }
}

impl PartialOrd for Edge {
    fn partial_cmp(&self, other: &Edge) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

fn read_line() -> String {
    let stdin = io::stdin();
    let mut buf = String::new();
    stdin.read_line(&mut buf).unwrap();
    buf
}

fn read_values<T>() -> Vec<T>
    where T: std::str::FromStr,
          T::Err: std::fmt::Debug
{
    read_line()
        .split(' ')
        .map(|a| a.trim().parse().unwrap())
        .collect()
}

会津合宿 2017 2 日目 G : Picnic

問題

http://judge.u-aizu.ac.jp/onlinejudge/cdescription.jsp?cid=ACPC2017Day2&pid=G

解法

ワーシャルフロイドと巡回セールスマン問題の bit DP で、ある町の部分集合を回って戻ってくるのにかかるコストを前計算しておく。

さらに、町の集合を の 2 つに分け、各集合のうちの部分集合を回った時に得られる幸福度を、個数制約付きナップサック問題をといて求める。

すると、 の部分集合と の部分集合を組み合わせることで、ある部分集合の最大値を求めることができる。

コード

import java.util.Scanner

object Main {
  def main(args: Array[String]): Unit = {
    val in = new Scanner(System.in)
    val N = in.nextInt()
    val X = in.nextInt()
    val Y = in.nextInt()
    val foods = (for (_ <- 0 until N) yield {
      val K = in.nextInt()
      (for (_ <- 0 until K) yield {
        val a = in.nextInt()
        val b = in.nextInt()
        val c = in.nextInt()
        (a, b, c)
      }).toArray
    }).toArray

    val dist = (for (_ <- 0 until N) yield (for (_ <- 0 until N) yield in.nextInt()).toArray).toArray
    for {
      k <- 0 until N
      i <- 0 until N
      j <- 0 until N
    } dist(i)(j) = math.min(dist(i)(j), dist(i)(k) + dist(k)(j))

    val cost = Array.fill(1 << (N + 1), N + 1)(Int.MaxValue / 2)
    cost(1)(0) = 0
    for {
      mask <- 0 until (1 << N)
      current <- 0 until N
      next <- 1 until (N + 1)
      if (mask & (1 << current)) != 0
      if (mask & (1 << next)) == 0
    } {
      val nextMask = mask | (1 << next)
      cost(nextMask)(next) = math.min(cost(nextMask)(next), cost(mask)(current) + dist(current)(next % N))
    }

    val bottom = N / 2
    val top = N - bottom

    val deq = new Array[Int](1000000)
    val deqv = new Array[Int](1000000)

    def knapsack(mask: Int, dp: Array[Int]): Unit = {
      dp(0) = 0
      for {
        i <- 0 until N
        if (mask & (1 << i)) != 0
        (price, value, remain) <- foods(i)
        a <- 0 until price
      } {
        var s = 0
        var t = 0

        var num = 0
        while (num * price <= (Y - a)) {
          val v = dp(num * price + a) - num * value
          while (s < t && deqv(t - 1) <= v) {
            t -= 1
          }
          deq(t) = num
          deqv(t) = v
          t += 1

          dp(num * price + a) = deqv(s) + num * value
          if (deq(s) == num - remain) {
            s += 1
          }

          num += 1
        }
      }

      for {
        i <- 1 to Y
      } {
        dp(i) = math.max(dp(i), dp(i - 1))
      }
    }

    val bottomDp = Array.fill[Int](1 << bottom, Y + 1)(-1)
    val topDp = Array.fill[Int](1 << top, Y + 1)(-1)
    for (mask <- 0 until (1 << bottom)) {
      knapsack(mask, bottomDp(mask))
    }
    for (mask <- 0 until (1 << top)) {
      knapsack(mask << bottom, topDp(mask))
    }

    var max = 0
    for {
      mask <- 0 until (1 << N)
    } {
      val moveCost = cost(mask | (1 << N))(N)
      val budget = math.min(Y, X - moveCost)
      val bottomMask = mask & ((1 << bottom) - 1)
      val topMask = mask >> bottom

      if (budget > 0) {
        for {
          bottomBudget <- 0 to budget
        } {
          max = math.max(max, bottomDp(bottomMask)(bottomBudget) + topDp(topMask)(budget - bottomBudget))
        }
      }
    }
    println(max)
  }
}

会津合宿 2017 3 日目 E : Taiyaki-Master and Eater

問題

AIZU ONLINE JUDGE

解法

2次元のBITを貼る。

コード

import java.util.Scanner

import scala.collection.mutable.ArrayBuffer

object Main extends App {
  val in = new Scanner(System.in)
  val H = in.nextInt()
  val W = in.nextInt()
  val T = in.nextInt()
  val Q = in.nextInt()

  var queries = new ArrayBuffer[(Int, Int, Int, Int, Int, Int)]()
  for (_ <- 0 until Q) {
    val t = in.nextInt()
    val c = in.nextInt()
    val h = in.nextInt()
    val w = in.nextInt()
    val (h2, w2) = if (c == 2) {
      (in.nextInt(), in.nextInt())
    } else {
      (0, 0)
    }
    queries.append((t, c, h, w, h2, w2))
    if (c == 0) {
      // 焼き上がり
      queries.append((t + T, -1, h, w, h2, w2))
    }
  }
  queries = queries.sortBy(q => (q._1, q._2))

  val f1 = new Fenwick2D(H, W)
  val f2 = new Fenwick2D(H, W)
  queries.foreach { case (_, c, h, w, h2, w2) =>
    c match {
      case -1 =>
        // 1->2
        if (f1.get(h, w) != 0) {
          f1.update(h, w, 0)
          f2.update(h, w, 1)
        }
      case 0 =>
        // 0->1
        f1.update(h, w, 1)
      case 1 =>
        // 2->0
        f2.update(h, w, 0)
      case 2 =>
        // count
        val a = f2.sum(h, w, h2, w2)
        val b = f1.sum(h, w, h2, w2)
        println(s"$a $b")
    }
  }


}

/**
  * 1-indexed 2-dimensional Fenwick's tree
  *
  * @param h height
  * @param w width
  */
class Fenwick2D(h: Int, w: Int) {
  val N: Int = h + 1
  val M: Int = w + 1
  val data: Array[Array[Long]] = Array.fill[Long](N + 1, M + 1)(0)

  def add(x: Int, y: Int, v: Long): Unit = {
    var i = x
    while (i <= N) {
      var j = y
      while (j <= M) {
        data(i)(j) += v
        j += j & -j
      }
      i += i & -i
    }
  }

  def update(x: Int, y: Int, v: Long): Unit = add(x, y, v - sum(x, y, x, y))

  def get(x: Int, y: Int): Long = sum(x, y, x, y)

  def sum(x0: Int, y0: Int, x1: Int, y1: Int): Long = {

    def partialSum(x: Int, y: Int): Long = {
      var res = 0L
      var i = x
      while (i > 0) {
        var j = y
        while (j > 0) {
          res += data(i)(j)
          j -= j & -j
        }
        i -= i & -i
      }
      res
    }

    partialSum(x1, y1) - partialSum(x0 - 1, y1) - partialSum(x1, y0 - 1) + partialSum(x0 - 1, y0 - 1)
  }


}

AOJ 2828: Matryoshka Doll

問題

Matryoshka Doll | Aizu Online Judge

解法

取り込まれた人形のコストを 0 として、取り込まれていない人形のコストはそのまま結果に加えるように、最小費用流のグラフを作る。

コード

import java.util.Scanner

import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer

object Main extends App {
  val in = new Scanner(System.in)

  def solve(n: Int): Long = {
    val xyz = (for (_ <- 0 until n) yield {
      Array(in.nextInt(), in.nextInt(), in.nextInt()).sorted
    }).toArray


    val source = n * 2
    val sink = source + 1
    val minimumCostFlow = new MinimumCostFlow(sink + 1)
    for {
      i <- 0 until n
      j <- 0 until n
      if i != j
      if xyz(i)(0) < xyz(j)(0)
      if xyz(i)(1) < xyz(j)(1)
      if xyz(i)(2) < xyz(j)(2)
    } {
      minimumCostFlow.addEdge(i, n + j, 1, 0)
    }
    for (i <- 0 until n) {
      minimumCostFlow.addEdge(source, i, 1, 0)
      minimumCostFlow.addEdge(i, sink, 1, xyz(i)(0) * xyz(i)(1) * xyz(i)(2))
      minimumCostFlow.addEdge(i + n, sink, 1, 0)
    }
    minimumCostFlow.calculateCost(source, sink, n)
  }

  def loop(): Unit = {
    while (true) {
      val n = in.nextInt()
      if (n == 0) {
        return
      }
      val ans = solve(n)
      println(ans)
    }
  }

  loop()

}

class Edge(val to: Int, var cap: Long, val cost: Long, val rev: Int)

class MinimumCostFlow(V: Int) {
  val graph: Array[ArrayBuffer[Edge]] = (for (_ <- 0 until V) yield new ArrayBuffer[Edge]()).toArray
  val prevV = new Array[Int](V)
  val prevE = new Array[Int](V)

  def addEdge(from: Int, to: Int, cap: Long, cost: Long): Unit = {
    graph(from).append(new Edge(to, cap, cost, graph(to).size))
    graph(to).append(new Edge(from, 0, -cost, graph(from).size - 1))
  }

  def calculateCost(source: Int, sink: Int, flow: Long): Long = {
    val INF = Long.MaxValue / 2
    var cost = 0L
    var residue = flow
    val potential = new Array[Long](V)
    while (residue > 0) {
      val dist = Array.fill[Long](V)(INF)
      dist(source) = 0
      val queue = mutable.PriorityQueue.empty[(Long, Int)](implicitly[Ordering[(Long, Int)]].reverse)
      queue.enqueue((0, source))
      while (queue.nonEmpty) {
        val p = queue.dequeue()
        val v = p._2
        if (dist(v) >= p._1) {
          for (i <- graph(v).indices) {
            val e = graph(v)(i)
            val u = e.to
            if (e.cap > 0 && dist(u) > dist(v) + e.cost + potential(v) - potential(u)) {
              dist(u) = dist(v) + e.cost + potential(v) - potential(u)
              prevV(u) = v
              prevE(u) = i
              queue.enqueue((dist(u), u))
            }
          }
        }
      }
      if (dist(sink) == INF) {
        return -1
      }

      for (v <- 0 until V) {
        potential(v) += dist(v)
      }
      var d = residue
      var v = sink
      while (v != source) {
        d = math.min(d, graph(prevV(v))(prevE(v)).cap)
        v = prevV(v)
      }
      residue -= d
      cost += d * potential(sink)

      v = sink
      while (v != source) {
        val e = graph(prevV(v))(prevE(v))
        e.cap -= d
        graph(v)(e.rev).cap += d
        v = prevV(v)
      }
    }
    return cost
  }
}

ARC082 F - Sandglass

問題

https://beta.atcoder.jp/contests/arc082/tasks/arc082_d

解法

解説動画の通りにやった。

最初に a 入っている時の t 秒後の砂の量は以下のように書ける。

この をシミュレーションしていけば良い。

コード

use std::io;
use std::str;
use std::usize;
use std::cmp;


static INF: i32 = 2000000000;

fn main() {
    let x = read_values::<i32>()[0];
    let k = read_values::<usize>()[0];
    let r = {
        let mut r = read_values::<i32>();
        r.push(INF);
        r
    };
    let q = read_values::<usize>()[0];
    let (t, a) = {
        let mut t = vec![INF; q + 1];
        let mut a = vec![INF; q + 1];
        for i in 0..q {
            let v = read_values::<i32>();
            t[i] = v[0];
            a[i] = v[1];
        }
        (t, a)
    };

    let mut y = 0;
    let mut x1 = 0;
    let mut x2 = x;
    let mut upper = true;

    let mut flat = false;
    let mut flat_height = 0;

    let mut r_head = 0;
    let mut t_head = 0;
    let mut prev_time = 0;
    while r_head < k || t_head < q {
        if r[r_head] < t[t_head] {
            // turn
            let time = r[r_head] - prev_time;
            if upper {
                y -= time;
                if y < 0 && x1 < -y {
                    x1 = -y;
                }
            } else {
                y += time;
                if y > 0 && x2 > x - y {
                    x2 = x - y;
                }
            }

            if flat {
                if upper {
                    flat_height -= time;
                } else {
                    flat_height += time;
                }
                flat_height = cmp::max(flat_height, 0);
                flat_height = cmp::min(flat_height, x);
            } else if x1 >= x2 {
                flat = true;
                if upper {
                    flat_height = 0;
                } else {
                    flat_height = x;
                }
            }

            prev_time = r[r_head];
            upper = !upper;
            r_head += 1;
        } else {
            // query
            let time = t[t_head] - prev_time;

            let mut ans = if a[t_head] < x1 {
                x1 + y
            } else if x2 < a[t_head] {
                x2 + y
            } else {
                a[t_head] + y
            };

            if upper {
                ans -= time;
            } else {
                ans += time;
            }

            if flat {
                ans = flat_height;
                if upper {
                    ans -= time;
                } else {
                    ans += time;
                }
            }

            ans = cmp::max(0, ans);
            ans = cmp::min(x, ans);
            println!("{}", ans);

            t_head += 1;
        }
    }
}

fn read_line() -> String {
    let stdin = io::stdin();
    let mut buf = String::new();
    stdin.read_line(&mut buf).unwrap();
    buf
}

fn read_values<T>() -> Vec<T>
    where T: std::str::FromStr,
          T::Err: std::fmt::Debug
{
    read_line()
        .split(' ')
        .map(|a| a.trim().parse().unwrap())
        .collect()
}