DP
ははーん、これはあれですね。
$cnt[0\leq k \leq 3][0 \leq r \leq 6]$ を、「カードを $k$ 枚選んでその和を7で割った余りが $r$ になる組み合わせの個数」としてDPして、最後に $cnt[3][0]$ を返せばいい感じですね？
カードを1枚も見ていない初期状態では $cnt[0][0] = 1, cnt[k][r] = 0$ として、カード $a_i$ に対して $cnt[k][r]$ を $cnt[k+1][(r + a_i) \bmod 7]$ に足し込めばできて、これは28要素の配列の21要素を更新するので、immutableな配列のままやってもオーバーヘッドは気にならないはず。
import Data.Array
solve :: Int -- N
-> [Int] -- a_i
-> Int -- 答え
solve _n as = stN ! (3,0)
where
st0 = accumArray (+) 0 ((0,0),(3,6)) [((0,0),1)]
stN = foldl step st0 as
step st a = accum (+) st
[ ((succ k, j), st ! (k, i))
| (i, j) <- [(i, mod (a + i) 7) | i <- [0 .. 6]]
, k <- [0 .. 2] ]
できました！
公式解説
「カードの内訳が、7で割って$r$余るようなカードはそれぞれ何枚なのかだけ数えて、それらから3枚選んで7で割り切れるようにする場合の数を数えるだけだよ」
DPいらんかった。
solve :: Int -> [Int] -> Int
solve _n as = flip div 6 $ sum
[ nI * nJ * nK
| (i, nI) <- assocs cnt , nI > 0, let cntI = accum (+) cnt [(i, -1)]
, (j, nJ) <- assocs cntI, nJ > 0, let cntJ = accum (+) cntI [(j, -1)]
, let ij = mod (i + j) 7, let k = if ij == 0 then 0 else 7 - ij
, let nK = cntJ ! k, nK > 0
]
where
cnt = accumArray (+) 0 (0,6) [(mod a 7, 1) | a <- as]
調べる場合の数はたった $7^3$ 通りだし、カードの枚数配列は要素数たったの7なので、immutableな更新を平気で使う。
カードの選び方に順番は関係ないので、順番を気にせず数え上げた結果を $_3 P_3 = 6$ で割っている。
公式解説だと、同じ枠から選んだ場合は、とかを $_n C_r$ を駆使して数えていたり、事前計算したリストを使ったりしてるけど、単純にこれでいける。
力こそパワー解
こちらの記事に、どストレートな解が。
「何も考えずにカードリストから3枚選ぶやり方を全て試して、7で割り切れるものを数えるよ」
$_{10^5} C_3 = 166,661,666,700,000$ 通りの数え上げは無理じゃね？と思いつつ、追試。
（元記事だと最初に mod 7 を計算しているけど、それは不要）
solve :: Int -> [Int] -> Int
solve _n as = length $ filter (\x -> 0 == mod (sum x) 7) $ combinations 3 as
combinations :: Int -> [a] -> [[a]]
combinations 0 _ = [[]]
combinations _ [] = []
combinations n (x:xs) = [x:ys | ys <- combinations (pred n) xs] ++ combinations n xs
間に合っちゃいました。何が起きてるの？
問題文には $N \leq 100000$ と書いてあるけど、そんなテストケースはないとかそんな感じ？
終わりに
スキルチェックでHaskell使えるようにしてホスイ…