HaskellでABC439を解く

はじめに

2026年の初の参加は残念ながら惨敗でした。

A問題

数式通りの計算をしてAC。

main = do
  n <- readLn :: IO Int
  print $ 2^n - 2*n

A問題提出

B問題

処理を進めて一桁に達した時に1の場合はTrueとしました。
厳密にはすでに処理した値と同じものが出てきた段階で打ち切るような処理が必要と思いましたが
時間をかけすぎるのも何なので嘘回答だと思いつつ提出したらACしました。いいのかな？

main = do
  ns <- map digitToInt <$> getLine
  putStrLn $ bool "No" "Yes" $ solve ns

solve :: [Int] -> Bool
solve [x] = x == 1 
solve xxs@(x:xs) = solve xxs1
  where
    val = sum $ map (^2) xxs
    xxs1 = map digitToInt $ show val

solve _ = error "error"

B問題提出

C問題

xの範囲は√nなのでO(3000)程度に収まっているので後処理部分がO(n)程度で処理できるので
間に合うはずと思って提出しましたがTLEを解消できずでした。

main = do
  n <- readLn :: IO Int
  let m = floor $ sqrt $ fromIntegral n -- 上限
  let cands = [(x,y) | x <- [1..m], let y = floor $ sqrt $ fromIntegral (n-x^2), x^2+y^2 <= n ]
  let ans = map snd $ filter ((==1).fst)$ map (\g->(length g, head g)) $ group $ sort $ solve cands
  print $ length ans
  putStrLn $ unwords $ map show ans

solve [] = []
solve ((x,y):xys)
  | x >= y = []
  | otherwise = ks ++ solve xys
  where
    ks = [x2+y'^2| y' <- [x+1..y], let x2=x^2]

C問題提出(TLE)

コンテスト後にAIにGoLangに変換してもらったらあっさりACしました。
アルゴリズムは問題ないということか。。。

package main

import (
	"bufio"
	"fmt"
	"math"
	"os"
	"sort"
)

func main() {
	in := bufio.NewReader(os.Stdin)

	var n int
	fmt.Fscan(in, &n)

	// m = floor(sqrt(n))
	m := int(math.Floor(math.Sqrt(float64(n))))

	// cands = [(x,y)]
	type Pair struct {
		x int
		y int
	}
	var cands []Pair
	for x := 1; x <= m; x++ {
		y := int(math.Floor(math.Sqrt(float64(n - x*x))))
		cands = append(cands, Pair{x, y})
	}

	// solve cands
	var vals []int
	for _, p := range cands {
		x := p.x
		y := p.y
		if x >= y {
			continue
		}
		x2 := x * x
		for yp := x + 1; yp <= y; yp++ {
			vals = append(vals, x2+yp*yp)
		}
	}

	// sort
	sort.Ints(vals)

	// group & count == 1
	var ans []int
	for i := 0; i < len(vals); {
		j := i + 1
		for j < len(vals) && vals[j] == vals[i] {
			j++
		}
		if j-i == 1 {
			ans = append(ans, vals[i])
		}
		i = j
	}

	// output
	out := bufio.NewWriter(os.Stdout)
	defer out.Flush()

	fmt.Fprintln(out, len(ans))
	for i, v := range ans {
		if i > 0 {
			fmt.Fprint(out, " ")
		}
		fmt.Fprint(out, v)
	}
	fmt.Fprintln(out)
}

C問題提出(AC:コンテスト後 GoLang)

sortしたりgroupとったりする辺りをaccumArrayで一気にやるようにしたら無事AC。

import qualified Data.Array.Unboxed as AU

main = do
  n <- readLn :: IO Int
  let m = floor $ sqrt $ fromIntegral n -- 上限
  let cands = [(x,y) | x <- [1..m], let y = floor $ sqrt $ fromIntegral (n-x^2)] :: [(Int,Int)]
  let cands2 = AU.accumArray (+) 0 (1,n) $ map (,1) $ solve cands :: AU.UArray Int Int
  let ans = [i | (i,1) <- AU.assocs cands2]
  print $ length ans
  putStrLn $ unwords $ map show ans

solve [] = []
solve ((x,y):xys)
  | x >= y = []
  | otherwise = ks ++ solve xys
  where
    ks = [x2+y'^2| y' <- [x+1..y], let x2=x^2]

C問題提出(AC:コンテスト後)

おわりに

C問題ArrayとかVectorを使わないと間に合わない事が判明。覚えておこう。