My Weblog

Blog about programming and math

SPOJ Power Sums

SPOJ Power Sums accepted after couple of attempts. This problem is similar to SRM 397 SumOfPowers . This discussion on topcoder forum shows couple of ideas to solve this problem using Gaussian elimination , recursive sum and Bernoulli numbers . First i implemented Gaussian elimination mod prime . The idea is
\sum_{i=1}^x i^k = a_{k+1}x^{k+1}+a_kx^k+....+a_1x+a_0 . For this equation , we have k+2 unknown coefficients so we need k+2 equations to solve this. Putting x from 0 to ( k + 1 ) and we can have k + 2 equation . Solve this system of equation using Gaussian elimination mod prime .
\begin{pmatrix} x_0^{k+1} & x_0^k & . & . & x_0 & 1 \\  x_1^{k+1} & x_1^k & . & . & x_1 & 1 \\  x_2^{k+1} & x_2^k & . & . & x_2 & 1 \\   . & . & . & . & . & . \\  . & . & . & . & . & . \\  x_{k+1}^{k+1} & x_{k+1}^k & . & . & x_{k+1} & 1\end{pmatrix} *    \begin{pmatrix} a_{k+1} \\  a_k \\  . \\   .  \\  a_1 \\  a_0 \end{pmatrix} = \begin{pmatrix} \sum_{i=1}^{x_0} i^k \\  \sum_{i=1}^{x_1} i^k \\  \sum_{i=1}^{x_2} i^k \\   .  \\  .\\  \sum_{i=1}^{x_{k+1}} i^k \end{pmatrix}{\mod p }
Haskell code for this implementation but it got time limit exceed . This code can also be useb for computation of determinant mod primes and combine the result using chinese remainder theorem.

import Data.List

extendedGcd :: Integer -> Integer -> ( Integer , Integer ) 
extendedGcd a b 
   | b == 0 = ( 1 , 0 ) 
   | otherwise = ( t , s - q * t ) where 
	( q , r ) = quotRem a b 
	( s , t ) = extendedGcd b r 
	 

modInv :: Integer -> Integer -> Integer 
modInv  a b 
   | gcd a b /= 1 = error " gcd is not 1 "
   | otherwise = d where 
	d = until ( > 0 ) ( + b  ) . fst.extendedGcd a $ b  



adjustwitH::[ Integer ] -> [ Integer ] -> [ Integer ]
adjustwitH ( m:ms ) ( n:ns ) = map ( flip mod 10007 )  $  zipWith ( - ) ms $ map ( ( m * modInv n 10007 )* )  ns --m/n mod p modular inverse



echeloN::[[ Integer ]]->[[ Integer ]]
echeloN rs 
    | null rs || null ( head rs) = rs
    | null rs_2 = map (0:) ( echeloN $ map tail rs )  --haskell road to logic 
    | otherwise = piv : map ( 0:) ( echeloN rs' )  where 
	rs'= map ( adjustwitH piv ) ( rs_1 ++ rs_3 )
	( rs_1 , rs_2 ) = span ( \( x:xs )-> x == 0 ) rs
	( piv:rs_3 ) = rs_2



rows , cols :: [[Integer ]] -> Int
rows m = length m 
cols m = length.head $ m 



backS :: [[Integer]] -> [Integer]
backS rs = backS' rs [] where 
	backS' [] ps = ps
	backS' rs ps = backS' rs' ( p:ps ) where 
	  a = ( last rs ) !! ( ( cols rs ) - 2 )
	  c = ( last rs ) !! ( ( cols rs ) - 1 )
	  p = mod ( c * modInv a 10007 ) 10007
	  rs' = eliminate p ( init rs ) 



eliminate :: Integer -> [[Integer]] -> [[Integer]]
eliminate p rs = map ( simplify p ) rs where 
	simplify p row = init ( init row ) ++ [ mod ( d - b * p) 10007 ] where 
		d = last row 
		b = last ( init row ) 


helpMod::Integer->Integer->Integer->Integer
helpMod a d n= fun a d where
	fun a 0 = 1 
	fun a 1 = mod a n
	fun a 2 = mod (a^2) n
	fun a d = let
			p=fun (mod (a^2) n) (div d 2)
			q=if odd d then mod (a*p) n else p
		  in mod q n

modS :: Integer -> Integer -> Integer
modS p k = mod ( sum [ helpMod i k 10007 | i <- [1..p]] )  10007

solve :: Integer -> [Integer]
solve k =  backS . echeloN $  [ [ helpMod i n 10007  | n <- [ ( k + 1 ) , k .. 0 ] ] ++ [ modS i k ] | i <- [1 .. (k+2)] ]

format ::Integer -> [Integer] -> String
format k [x]  = show x ++ "x^" ++ show k 
format k t@(x:xs)  = ( show x ++ "x^" ++ show k ++ " + ") ++ format  ( k - 1 ) xs 

final :: Integer -> String
final k = format ( k + 1) ( solve k )

main = interact $ unlines. map ( final.read ) . lines 

My second attempt was recursive formula ( m + 1 )\sum_{k=0}^m = ( n + 1 ) ^{m+1} -  \sum_{r=2}^{m+1} ( {{n}\choose{r}} *\sum_{k=0}^n k^{m+1-r} ) but unfortunately this one also time limit exceed .

import Data.List
import Data.Maybe
import qualified Data.ByteString.Char8 as BS 


polySum :: [Integer] -> [Integer] -> [Integer]
polySum ( x : xs ) ys =map ( flip mod 10007 ) $  x : zipWith ( + ) xs ys 

polySub :: [Integer] -> [Integer] -> [Integer]
polySub (x:xs) ys = map ( flip mod 10007 ) $ x : zipWith ( - ) xs ys 


binomM :: Integer -> Integer -> Integer -> Integer
binomM n k p =  pas !! ( fromIntegral n) !! ( fromIntegral k )   



pas :: [[Integer]]
pas = iterate ( \x -> map ( flip mod 10007 ) $  1 : zipWith (+) x ( tail x )   ++ [1] ) [1] 


extendedGcd :: Integer -> Integer -> ( Integer , Integer ) 
extendedGcd a b  
   | b == 0 = ( 1 , 0 ) 
   | otherwise = ( t , s - q * t ) where 
	( q , r ) = quotRem a b 
	( s , t ) = extendedGcd b r 
	 

modInv :: Integer -> Integer -> Integer 
modInv  a b 
   | gcd a  b /= 1 = error " gcd is not 1 "
   | otherwise = d where 
	d = until ( > 0 ) ( + b  ) . fst.extendedGcd a  $ b   


sum' ::[[Integer]] -> [Integer]
sum' xs = foldr ( \x y -> polySum x y ) [0] xs 

solve :: Integer -> [[Integer]]
solve m = foldl ( \xs y -> ( polySub ( [ modInv ( y + 1) 10007 * binomM ( y + 1 ) i 10007 | i <- [ 0 .. ( y + 1 ) ] ] ) ( sum' $ zipWith  map    [ ( ( modInv ( y + 1 ) 10007 * binomM ( y + 1 ) r  10007 )* ) | r <- [ 2 .. ( y + 1 ) ] ] xs  ) ) : xs  ) [[1,1]] [ 1 .. m ]

Finally attempt was using Bernoulli numbers and bingo. Sum is given by S_m(n) = {1\over{m+1}}\sum_{k=0}^m {m+1\choose{k}} B_k n^{m+1-k} and Bernoulli numbers can be computed using B_m= 1 - \sum_{k=0}^{m-1}\binom mk\frac{B_k}{m-k+1} .Accepted Haskell code . Learn lot while solving this problem.

import Data.List
import Data.Maybe
import qualified Data.ByteString.Char8 as BS


readInt :: BS.ByteString -> Integer
readInt = fst . fromJust . BS.readInteger 


binomM :: Integer -> Integer -> Integer -> Integer
binomM n k p = pas !! ( fromIntegral n ) !! ( fromIntegral k ) 

pas :: [[Integer]]
pas = iterate ( \x ->  1 : zipWith (\ a b -> mod ( a + b ) 10007 ) x ( tail x ) ++ [1] ) [1] 


extendedGcd :: Integer -> Integer -> ( Integer , Integer ) 
extendedGcd a b 
   | b == 0 = ( 1 , 0 ) 
   | otherwise = ( t , s - q * t ) where 
	( q , r ) = quotRem a b 
	( s , t ) = extendedGcd b r 
	 

modInv :: Integer -> Integer -> Integer 
modInv  a b 
   | gcd a b /= 1 = error " gcd is not 1 "
   | otherwise = d where 
	d = until ( > 0 ) ( + b  ) . fst.extendedGcd a $ b   


ber ::  [Integer]
ber  = foldl ( \ xs y ->  if and [ odd y , y /= 1 ] then xs ++ [0] else  xs ++ [ mod ( 1 -  ( sum $ zipWith (*) [  binomM y k 10007   * modInv ( y - k + 1 ) 10007 | k <-[ 0 .. pred y ]  ]  xs ) ) 10007 ]  ) [1] [1..300]

{--ber is fast than berIt 
berIt :: [( [Integer] , Integer )]
berIt = iterate ( \( xs , y) -> ( map ( flip mod 10007 ) $ if and [ odd y , y /= 1 ] then xs ++ [0] else  xs ++ [ 1 -  ( sum $ zipWith (*) [   binomM y k 10007   * modInv ( y - k + 1 ) 10007 | k <-[ 0 .. pred y ]  ]  xs ) ]  , y + 1  ) )  ( [1] , 1 ) 
 --}

solve :: Integer -> [Integer]
solve m =   [ mod ( binomM ( m + 1 ) k 10007 * ber !! ( fromIntegral k )  * modInv ( m + 1 ) 10007 ) 10007  | k <- [0..m] ] 


final :: Integer -> BS.ByteString 
final k = format ( k + 1 ) ( solve k )

format::Integer -> [Integer] -> BS.ByteString
format _ [] = BS.pack "0x^0"
format k ( x : xs ) = BS.append ( BS.append ( BS.append ( BS.append ( BS.pack.show $ x ) ( BS.pack  "x^" ) ) ( BS.pack.show $ k ) ) ( BS.pack " + " ) ) ( format ( k - 1 ) xs  )


main = BS.interact $ BS.unlines . map ( final . readInt )  . BS.lines  

Here is more faster way to generate Bernoulli numbers. Replace ber !! ( fromIntegral k ) in solve m with bernoulli k for improved run time.

bernoulli :: Integer -> Integer
bernoulli = genericIndex bernoulli_table
bernoulli_table = map bernoulli' [0..]

bernoulli' :: Integer ->  Integer
bernoulli' 0 = 1
bernoulli' 1 = 1 * modInv 2 10007
bernoulli' n | odd n          = 0
             | n `mod` 6 == 0 = let val = (n+3) * modInv 3 10007 - bernoulliA n (n `div` 6) in
                                  mod ( val * modInv ( binomM ( n + 3 ) n 10007 ) 10007 ) 10007 --fromIntegral (choose (n+3) n)
             | n `mod` 6 == 2 = let val = (n+3) * modInv 3 10007 - bernoulliA n ((n-2) `div` 6) in
                                  mod ( val * modInv  ( binomM ( n + 3 ) n 10007 ) 10007 ) 10007 --fromIntegral (choose (n+3) n)
             | n `mod` 6 == 4 = let val = -(n+3) * modInv 6 10007 - bernoulliA n ((n-4) `div` 6) in
                                 mod (  val * modInv ( binomM ( n + 3 ) n 10007 ) 10007 ) 10007 -- fromIntegral (choose (n+3) n)



bernoulliA :: Integer -> Integer ->  Integer
bernoulliA n m = mod ( sum (map foo [1..m]) ) 10007
  where foo j = binomM  (n+3) (n-6*j) 10007  *  bernoulli (n-6*j)
Advertisements

July 15, 2011 - Posted by | Programming | , , , ,

No comments yet.

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s

%d bloggers like this: