1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

import Data.Maybe (mapMaybe)
import Number (isPrime')

t :: (Integral a) => a -> a -> Maybe a
t n i | n < i || i < 1 = Nothing
      | otherwise = Just $ div (n*(n-1)) 2 + i

isPrime :: (Integer, Integer) -> Bool
isPrime = maybe False isPrime'.uncurry t

isTriple :: Integer -> Integer -> Bool
isTriple n i
    | not $ isPrime (n,i) = False
    | null ps              = False
    | length ps > 1        = True
    | otherwise            = any isPrime.nns.head $ ps
    where ns | even n    = [(n-1,i-1),(n-1,i+1),(n+1,i)]
             | otherwise = [(n-1,i),(n+1,i+1),(n+1,i-1)]
          ps = filter isPrime ns
          nns (m,j)
              | (m,j) == (n-1,i)   = [(n-2,i-1),(n-2,i+1)]
              | (m,j) == (n+1,i+1) = [(n,i+2),(n+2,i+1)]
              | (m,j) == (n+1,i-1) = [(n+2,i-1),(n,i-2)]
              | (m,j) == (n-1,i-1) = [(n,i-2),(n-2,i-1)]
              | (m,j) == (n-1,i+1) = [(n-2,i+1),(n,i+2)]
              | (m,j) == (n+1,i)   = [(n+2,i-1),(n+2,i+1)]

p196 :: Integer -> Integer
p196 n = sum.mapMaybe (t n) . filter (isTriple n) $ [1..n]

main = print $ p196 5678027 + p196 7208785

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66

import static java.lang.Math.*;
import java.util.BitSet;
 
public class P196{
    static int a = 5678027, b = 7208785, up = (int)(sqrt(triangle(b+2, b+2))/2); //sieve range
    static BitSet prime,prime_part;
  
    public static void main(String[] args){
        prime = new BitSet(up); // prime(i) is true iff 2i+3 is prime
        prime.flip(0, up); // defalt is true
        for(int i = 0; i >= 0; i = prime.nextSetBit(i+1)) // sieve prime
            for(long j = (2*(long)i+3)*(2*(long)i+3); j < 2*up+3 ; j+=4*i+6)
                prime.clear((int)((j-3)/2));

        System.out.println(solve(a)+solve(b));
    }
    
    // find S(n)
    public static long solve(int n){
        long s = triangle((long)n-2,1), t = triangle((long)n+2, (long)n+2), c = 0;
        sieve(s,t);
        for(int i = 1; i <= n; i++)
            if(isTriple(s, t, n, i))
                c+= triangle(n, i);
        return c;
    }

    public static boolean isTriple(long s, long t, int n, int i){
        int[] n1 = {1,-1,-1}, i1 = {0,-1,1}; // neighbors 1
        int[][] n2 = {{2,2},{0,-2},{-2,0}}, i2 = {{1,-1},{-2,-1},{1,2}}; // neighbors 2
        int p = 0, q = -1, k = n%2 == 0? 1: -1;
        
        if(!isPrime(s, t, triangle(n,i))) return false;
        for(int j = 0; j < 3; j++) // neighbors check 1
            if(isPrime(s, t, triangle(n+k*n1[j], i+k*i1[j]))){
                p ++; q = j;
                if (p > 1) return true;
            }
        if(p == 0) return false;
         for(int j = 0; j < 2; j++) // neighbors check 2
             if(isPrime(s, t, triangle(n+k*n2[q][j], i+k*i2[q][j])))
                 return true;
        return false;
    }
    // number of row n, col i    
    public static long triangle(long n, long i){
        if(n < i || i < 1) return 0;
        return n*(n-1)/2+i;
    }

    // sieve prime in [s,t]
    public static void sieve(long s, long t){
        int size = (int)(t/2 - s/2);
        prime_part = new BitSet(size); // odd i is prime iff prime((i-s)/2) is true
        prime_part.flip(0, size); // default true
        for(int i = 0; i >=0 && 2*i+3 <= t; i = prime.nextSetBit(i+1)){
            long p = 2*i+3, k = s-p*p <= 0? 0: (s-p*p)/(2*p)+1; // s =< j
            for(long j = (p+2*k)*p; j <= t ; j+=2*p) // sieve
                prime_part.clear((int)((j-s)/2));
        }
    }
    public static boolean isPrime(long s, long t, long q){
        if(q % 2 == 0) return false;
        return prime_part.get((int)(q-s)/2);
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81

{-# OPTIONS_GHC -XBangPatterns -XPatternGuards#-}

import Data.Array.ST (STUArray, newArray, readArray, writeArray)
import Control.Monad (when, liftM, filterM)
import Control.Monad.ST (ST, runST)

infixl 9 .->
(.->) :: Monad m => m Bool -> m a -> m a -> m a
(.->) b v1 v2 = do p <- b; if p then v1 else v2                           

whenM :: Monad m => m Bool -> m () -> m ()
whenM b v = do p <- b; if p then v else return ()

primesUpTo :: Integer -> ST s (STUArray s Integer Bool)
primesUpTo n =
    do p <- newArray (1, div n 2) True
       let u = floor.sqrt.fromIntegral $ n
           loop !i =
               when (i<=u) $
               do whenM (readArray p $ div i 2) $ sieve (i*i) (2*i) 
                  loop (i+2)
           sieve !i !j =
               when (i<=n) $
               do writeArray p (div i 2) False
                  sieve (i+j) j
       loop 3
       return p 

primesRange :: STUArray s Integer Bool -> Integer -> Integer -> ST s (STUArray s Integer Bool)
primesRange p s t =
    do q <- newArray (div s 2, div t 2) True
       let u = floor.sqrt.fromIntegral $ t
           loop !i =
               when (i<=u) $
               do whenM (readArray p $ div i 2) $
                    let k = if s > i*i then 1 + div (s-i*i) (2*i)
                            else 0
                    in sieve ((i+2*k)*i) (2*i)
                  loop (i+2)
           sieve !i !j =
               when (i<=t) $
               do writeArray q (div i 2) False
                  sieve (i+j) j
       loop 3
       return q

t :: (Integral a) => a -> a -> a
t n i = div (n*(n-1)) 2 + i

p196 :: STUArray s Integer Bool -> Integer -> ST s Integer
p196 p n =
    do q <- primesRange p (t (n-2) 1) (t (n+2) (n+2))
       let isPrime (m,j) | m < j || j < 1 = return False
                         | even $ t m j   = return False
                         | otherwise      = readArray q $ t m j `div` 2 
           isTriple i = liftM not (isPrime (n,i)) .-> return False $
                        liftM null ps             .-> return False $
                        liftM ((>1).length) ps    .-> return True  $
                        liftM or.mapM isPrime' =<< liftM (nns.head) ps
               where k = if even n then 1 else -1
                     isPrime' (m,j) = isPrime (n+k*m, i+k*j)
                     ns = [(1,0),(-1,-1),(-1,1)]
                     ps = filterM isPrime' ns
                     nns (m,j)
                      | (m,j) == ns!!0 = [(2,1),(2,-1)]
                      | (m,j) == ns!!1 = [(0,-2),(-2,-1)]
                      | (m,j) == ns!!2 = [(-2,1),(0,2)]
           loop !i !c
               | i > n     = return c
               | otherwise = isTriple i .-> loop (i+2) (c+t n i) $
                             loop (i+1) c
       loop 1 0

a = 5678027
b = 7208785

main = print.runST $
       do p <- primesUpTo.floor.sqrt.fromIntegral $ t (b+2) (b+2)
          sa <- p196 p a
          sb <- p196 p b
          return $ sa + sb