新年一発目のAtCoder振り返り。
昨年までnoteに振り返りを書いていましたが、noteは技術記事より別のことに使いたいなと思い今年からQiitaで書くことにしました。
前回の記事はこちら
https://note.com/flower742/n/n5e1221333d8c
今回の結果
A,Bまでは特に問題なかったのですが、CのTLEが解消できずに2完で終わってしまいました。
各問題の振り返り
A - 2023
末尾の文字列を4に置換する問題です。
reverseとtailを使って末尾を取り除いて4を先頭に結合し、再度reverseで元に戻して出力しました。
main :: IO ()
main = do
s <- getLine
putStrLn $ reverse $ "4" ++ (tail . reverse $ s)
B - Tetrahedral Number
x+y+z≤Nを満たす非負整数の組 (x,y,z)を辞書順に出力する問題です。
制約を見ると0≤N≤21かつNが整数であることから、全探索で問題なさそうと判断しました。
こういう時はリスト内包表記が便利ですね。
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
C - Loong Tracking
1からN までの番号がついたN個のパーツからなる線を龍と呼び、龍の二次元座標上の移動を計算する問題です。
龍のパーツは最初(1, 0), (2, 0), (3, 0) .. (N, 0)にあり、このうち(1,0)にあるパーツを頭としています。
コンテストではこのように、queryをパースして龍の座標を更新することで解こうとしました。
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)
この場合、龍の位置の更新は
VU.cons newHead $ VU.init dragon
でやっているため、龍の長さをNとすると計算量がO(N)になります。
そして、これをqueryに対してforM_するので、queryの長さをQとすると全体の計算量はO(QN)になります。
今回の制約条件が
2≤N≤10^6,
1≤Q≤2×10^5
であることから、計算に時間がかかりTLEになってしまいました。
コンテスト終了後他の回答を見てみたのですが、みたところ2つの回答がありそうです。
-
scanlであらかじめ龍の軌道を計算しておくことで、計算量を抑える。この場合計算量はO(N+Q)になる。 - 龍の位置を
Seqで管理する。先頭への追加はO(1)で、位置の取得はO(log(min(i, N-i)))なので、計算量はO(N + Q * log(N))となる。
2.のupsolvedした回答は以下の通りです。
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した回答は以下の通りです。
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に限らないのでしょうが、速度を意識すると関数の計算量を頭に叩き込んだり色々工夫をしないといけないので難しいですね。
慣れが必要なところもあると思うので、引き続き精進して行きます。