AtCoder Regular Contest 089 D: Checker

AtCorderを題材にGoを勉強してみようと思い、ついでにBlog開設。

Scalaで書いてGoに翻訳する感じで進めてみる。

問題

arc089.contest.atcoder.jp

コード（初版 Scala 2.11.7）

累積和使わずに普通に書くとTLE。

Immutable原理主義の限界。

https://arc089.contest.atcoder.jp/submissions/2482495

object Main extends App {
  val Array(n, k) = io.StdIn.readLine().split(' ').map(_.toInt)
  val k2 = k*2
  val list = for (i <- 1 to n) yield {
    val Array(x, y, color) = io.StdIn.readLine().split(' ')
    // 0 ≦ x,y ≦ 2K に圧縮
    // 'W'はY方向にKずらした'B'と同値
    Array(x.toInt % k2, (y.toInt + (if (color == "B") 0 else k)) % k2)
  }
  val max = (for {
    i <- (0 until k)
    j <- (0 until k)
  } yield {
    val cnt = list.count { arr =>
      val x = (arr(0) + i) % k2
      val y = (arr(1) + j) % k2
      ((x < k && y < k) || (x >= k && y >= k))
    }
    Math.max(cnt, n - cnt)
  }).max
  print(max)
}

コード（累積和版 Scala 2.11.7）

累積和を事前算出することにより無事AC。

破壊的変更も状況によっては選択すべきという良い例かと思う。

https://arc089.contest.atcoder.jp/submissions/2482667

object Main extends App {
  val Array(n, k) = io.StdIn.readLine().split(' ').map(_.toInt)
  val k2 = k*2  
  // 0 ≦ x,y ≦ 2K に圧縮
  val squares = Array.fill(k2)(Array.fill(k2)(0))
  for (i <- 1 to n) {
    val Array(x, y, color) = io.StdIn.readLine().split(' ')
    // 'W'はY方向にKずらした'B'と同値
    squares(x.toInt % k2)((y.toInt + (if (color == "B") 0 else k)) % k2) += 1
  }
  // 計算量削減のために累積和を事前算出
  for {
    i <- 0 until k2
    j <- 1 until k2
  } squares(i)(j) += squares(i)(j-1)
  for {
    j <- 0 until k2
    i <- 1 until k2
  } squares(i)(j) += squares(i-1)(j)
  var max = 0
  val end = k2-1
  for {
    i <- 0 until k
    j <- 0 until k
  } {
    // ex) k = 2
    //      ■ □ □ ■ 
    //  j+k □ ■ ■ □
    //      □ ■ ■ □
    //    j ■ □ □ ■
    //      i   i+k
    val cnt =
      (squares(i)(j)) +
      (squares(end)(j) - squares(i+k)(j)) +
      (squares(i+k)(j+k) - squares(i+k)(j) - squares(i)(j+k) + squares(i)(j)) +
      (squares(i)(end) - squares(i)(j+k)) +
      (squares(end)(end) - squares(end)(j+k) - squares(i+k)(end) + squares(i+k)(j+k))
    max = Math.max(Math.max(cnt, n - cnt), max)
  }
  print(max)
}

コード（累積和版 Go 1.6）

翻訳というより機械的な置換で完成した。

爆速を期待したがデータが多いときは何故かScalaの方が速い結果となった。

https://arc089.contest.atcoder.jp/submissions/2482812

package main

import "fmt"

func main() {
  var n int
  var k int
  fmt.Scan(&n, &k)
  k2 := k*2
  // 0 ≦ x,y ≦ 2K に圧縮
  squares := make([][]int, k2)
  for i := range squares {
    squares[i] = make([]int, k2)
  }
  for i := 0; i < n; i++ {
    var x int
    var y int
    var color string
    fmt.Scan(&x, &y, &color)
    // 'W'はY方向にKずらした'B'と同値
    if color == "W" {
      y += k
    }
    squares[x % k2][y % k2]++ 
  }
  // 計算量削減のために累積和を事前算出
  for i := 0; i < k2; i++ {
    for j := 1; j < k2; j++ {
      squares[i][j] += squares[i][j-1]
    }      
  }
  for j := 0; j < k2; j++ {
    for i := 1; i < k2; i++ {
      squares[i][j] += squares[i-1][j]
    }
  }  
  max_ := 0
  end := k2-1
  for i := 0; i < k; i++ {
    for j := 0; j < k; j++ {
      // ex) k = 2
      //      ■ □ □ ■ 
      //  j+k □ ■ ■ □
      //      □ ■ ■ □
      //    j ■ □ □ ■
      //      i   i+k
      cnt :=
        (squares[i][j]) +
        (squares[end][j] - squares[i+k][j]) +
        (squares[i+k][j+k] - squares[i+k][j] - squares[i][j+k] + squares[i][j]) +
        (squares[i][end] - squares[i][j+k]) +
        (squares[end][end] - squares[end][j+k] - squares[i+k][end] + squares[i+k][j+k])
      max_ = max(max(cnt, n - cnt), max_)
    }
  }  
  fmt.Print(max_)
}

func max(a int, b int) int {
  if a > b {
    return a
  } else {
    return b
  }
}

コード（累積和 + fmt.Scan非使用版 Go 1.6）

qiita.com

@tnoda_さんの記事で競技プログラミング用の標準入力方法があったので組み込んでみた。

やはりGoは爆速だった。

https://arc089.contest.atcoder.jp/submissions/2482856

package main

import (
  "bufio"
  "fmt"
  "os"
  "strconv"
)

var sc = bufio.NewScanner(os.Stdin)

func nextInt() int {
    sc.Scan()
    i, e := strconv.Atoi(sc.Text())
    if e != nil {
        panic(e)
    }
    return i
}

func nextString() string {
  sc.Scan()
  return sc.Text()
}

func main() {
  var n int
  var k int
  fmt.Scan(&n, &k)
  k2 := k*2
  // 0 ≦ x,y ≦ 2K に圧縮
  squares := make([][]int, k2)
  for i := range squares {
    squares[i] = make([]int, k2)
  }
  sc.Split(bufio.ScanWords)
  for i := 0; i < n; i++ {
    x := nextInt()
    y := nextInt()
    color := nextString()
    // 'W'はY方向にKずらした'B'と同値
    if color == "W" {
      y += k
    }
    squares[x % k2][y % k2]++ 
  }
  // 計算量削減のために累積和を事前算出
  for i := 0; i < k2; i++ {
    for j := 1; j < k2; j++ {
      squares[i][j] += squares[i][j-1]
    }      
  }
  for j := 0; j < k2; j++ {
    for i := 1; i < k2; i++ {
      squares[i][j] += squares[i-1][j]
    }
  }  
  max_ := 0
  end := k2-1
  for i := 0; i < k; i++ {
    for j := 0; j < k; j++ {
      // ex) k = 2
      //      ■ □ □ ■ 
      //  j+k □ ■ ■ □
      //      □ ■ ■ □
      //    j ■ □ □ ■
      //      i   i+k
      cnt :=
        (squares[i][j]) +
        (squares[end][j] - squares[i+k][j]) +
        (squares[i+k][j+k] - squares[i+k][j] - squares[i][j+k] + squares[i][j]) +
        (squares[i][end] - squares[i][j+k]) +
        (squares[end][end] - squares[end][j+k] - squares[i+k][end] + squares[i+k][j+k])
      max_ = max(max(cnt, n - cnt), max_)
    }
  }  
  fmt.Print(max_)
}

func max(a int, b int) int {
  if a > b {
    return a
  } else {
    return b
  }
}

パフォーマンス比較

	処理時間（最大）	処理時間（最小）	メモリ使用量
Scala 初版	TLE	327 ms	-
Scala 累積和版	1,007 ms	320 ms	88,024 KB
Go 累積和版	1,556 ms	1 ms	45,312 KB
Go fmt.Scan非使用版	186 ms	1 ms	37,120 KB