AtCoder Beginner Contest 335 振り返り

新年一発目のAtCoder振り返り。
昨年までnoteに振り返りを書いていましたが、noteは技術記事より別のことに使いたいなと思い今年からQiitaで書くことにしました。

前回の記事はこちら
https://note.com/flower742/n/n5e1221333d8c

## 今回の結果

A,Bまでは特に問題なかったのですが、CのTLEが解消できずに2完で終わってしまいました。

![スクリーンショット 2024-01-11 9.47.49.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/167729/bf987665-3b58-9372-a2d0-6dcbdfa7b5d5.png)

![スクリーンショット 2024-01-11 9.42.30.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/167729/11ffdbc8-2c34-f532-9289-5b5bbb4a0bab.png)

## 各問題の振り返り

### A - 202
3


- [問題文](https://atcoder.jp/contests/abc335/tasks/abc335_a)
- [自分の回答](https://atcoder.jp/contests/abc335/submissions/49076390)

末尾の文字列を`4`に置換する問題です。
`reverse`と`tail`を使って末尾を取り除いて`4`を先頭に結合し、再度`reverse`で元に戻して出力しました。

```haskell
main :: IO ()
main = do
  s <- getLine

  putStrLn $ reverse $ "4" ++ (tail . reverse $ s)
```

### B - Tetrahedral Number

- [問題文](https://atcoder.jp/contests/abc335/tasks/abc335_b)
- [自分の回答](https://atcoder.jp/contests/abc335/submissions/49092424)

`x+y+z≤N`を満たす非負整数の組 `(x,y,z)`を辞書順に出力する問題です。
制約を見ると`0≤N≤21`かつNが整数であることから、全探索で問題なさそうと判断しました。　
こういう時はリスト内包表記が便利ですね。

```haskell
main :: IO ()
main = do
  n <- readInt

  let res = [(x, y, z) | x <- [0 .. n], y <- [0 .. n], z <- [0 .. n], x + y + z <= n]

  mapM_ (putStrLn . (\(x, y, z) -> show x ++ " " ++ show y ++ " " ++ show z)) res
```

- [問題文](https://atcoder.jp/contests/abc335/tasks/abc335_c)
- [自分の回答](https://atcoder.jp/contests/abc335/submissions/49092424)

### C - Loong Tracking

- [問題文](https://atcoder.jp/contests/abc335/tasks/abc335_c)
- [自分の回答(TLE)](https://atcoder.jp/contests/abc335/submissions/49117810)

 
1からN までの番号がついたN個のパーツからなる線を龍と呼び、龍の二次元座標上の移動を計算する問題です。
龍のパーツは最初`(1, 0), (2, 0), (3, 0) .. (N, 0)`にあり、このうち`(1,0)`にあるパーツを`頭`としています。

コンテストではこのように、queryをパースして龍の座標を更新することで解こうとしました。

```haskell
type Coord = (Int, Int)

type Dragon = VU.Vector Coord

type Query = (Int, String)

main :: IO ()
main = do
  [n, q] <- readInputIntList
  qs :: [Query] <- replicateM q $ do
    [m, r] <- words <$> getLine
    return (read m :: Int, r)

  let initCoord = VU.fromList [(i, 0) | i <- [1 .. n]]

  V.foldM_
    ( \acc (n, s) ->
        case n of
          1 -> return $ moveD s acc
          2 -> do
            let pos = getPos acc $ read @Int s
            putStrLn $ show (fst pos) ++ " " ++ show (snd pos)
            return acc
    )
    initCoord
    $ V.fromList qs

moveD :: String -> Dragon -> Dragon
moveD dir dragon =
  let dragonHead = VU.head dragon
      newHead = case dir of
        "R" -> (fst dragonHead + 1, snd dragonHead)
        "L" -> (fst dragonHead - 1, snd dragonHead)
        "U" -> (fst dragonHead, snd dragonHead + 1)
        "D" -> (fst dragonHead, snd dragonHead - 1)
        _ -> dragonHead
   in VU.cons newHead $ VU.init dragon

getPos :: Dragon -> Int -> Coord
getPos dragon p = dragon VU.! (p - 1)
```

この場合、龍の位置の更新は
```haskell
VU.cons newHead $ VU.init dragon
```

でやっているため、龍の長さを`N`とすると計算量が`O(N)`になります。
そして、これをqueryに対して`forM_`するので、queryの長さを`Q`とすると全体の計算量は`O(QN)`になります。

今回の制約条件が

```math
2≤N≤10^6, 
1≤Q≤2×10^5
```
であることから、計算に時間がかかりTLEになってしまいました。

コンテスト終了後他の回答を見てみたのですが、みたところ2つの回答がありそうです。

1. `scanl`であらかじめ龍の軌道を計算しておくことで、計算量を抑える。この場合計算量は`O(N+Q)`になる。
2. 龍の位置を`Seq`で管理する。先頭への追加は`O(1)`で、位置の取得は`O(log(min(i, N-i)))`なので、計算量は`O(N + Q * log(N))`となる。

2.のupsolvedした[回答](https://atcoder.jp/contests/abc335/submissions/49188771)は以下の通りです。

```haskell
type Query = (Int, String)

main :: IO ()
main = do
  [n, q] <- readInputIntList
  qs :: [Query] <- replicateM q $ do
    [i, d] <- words <$> getLine
    return (read i :: Int, d)

  let init = [(i, 0) | i <- [n, n - 1 .. 1]]

  let post = L.scanl' f (1, 0) [head d | (i, d) <- qs, i == 1]
        where
          f (i, j) 'L' = (i - 1, j)
          f (i, j) 'R' = (i + 1, j)
          f (i, j) 'U' = (i, j + 1)
          f (i, j) 'D' = (i, j - 1)
          f _ _ = error "invalid"

  let as = listArray @Array (-n + 1, q) $ init ++ tail post

  foldM_
    ( \acc q ->
        case q of
          (1, _) -> return (acc + 1)
          (2, p) -> do
            let pos = read @Int p
                (x, y) = as ! (acc - pos + 1)
            putStrLn $ show x ++ " " ++ show y
            return acc
          _ -> error "invalid"
    )
    0
    qs
```

2.のupsolvedした[回答](https://atcoder.jp/contests/abc335/submissions/49187840)は以下の通りです。

```haskell
type Query = (Int, String)

main :: IO ()
main = do
  [n, q] <- readInputIntList
  qs :: [Query] <- replicateM q $ do
    [i, d] <- words <$> getLine
    return (read i :: Int, d)

  let seq = Seq.fromList [(i, 0) | i <- [1 .. n]]
  let ans = catMaybes $ solve seq qs

  forM_ ans $ \(i, j) -> do
    putStrLn $ show i ++ " " ++ show j

solve :: (Eq a1, Num a1, Num a2, Num b) => Seq.Seq (a2, b) -> [(a1, String)] -> [Maybe (a2, b)]
solve seq [] = []
solve seq ((i, j) : ijs)
  | i == 1 && j == "R" = solve ((k + 1, l) Seq.<| seq) ijs
  | i == 1 && j == "L" = solve ((k - 1, l) Seq.<| seq) ijs
  | i == 1 && j == "U" = solve ((k, l + 1) Seq.<| seq) ijs
  | i == 1 && j == "D" = solve ((k, l - 1) Seq.<| seq) ijs
  | otherwise = Seq.lookup (p - 1) seq : solve seq ijs
  where
    ((k, l) Seq.:<| as) = seq
    p = read j :: Int
```

この回答だとTLEがなくなり無事ACになります。

## 全体を振り返って

Haskellに限らないのでしょうが、速度を意識すると関数の計算量を頭に叩き込んだり色々工夫をしないといけないので難しいですね。
慣れが必要なところもあると思うので、引き続き精進して行きます。