-
Notifications
You must be signed in to change notification settings - Fork 0
/
pairs.hs
53 lines (45 loc) · 1.34 KB
/
pairs.hs
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
import Data.List
import Data.Functor
import qualified Data.Vector as V
find_pairs k x ys = find_pairs_ k x ys
find_pairs_ :: Int -> Int -> V.Vector Int -> Int
find_pairs_ k x ys
| V.length ys < 1 = 0
| V.length ys == 1 =
let y = V.head ys
in if x - y == k then 1 else 0
| otherwise =
let half = (V.length ys) `div` 2
atfirst = V.head ys
athalf = (V.!) ys half
(left, right) = V.splitAt half ys
lpairs = V.takeWhile (\y -> x - y == k) ys
in
if x - atfirst == k
then V.length lpairs
else if x - athalf < k
then find_pairs k x left
else find_pairs k x right
isolve :: Int -> V.Vector Int -> V.Vector Int -> Int
isolve k arr rarr
| V.length arr < 2 = 0
| otherwise =
let
x = V.head arr
xs = V.tail arr
rx = V.last rarr
rxs = V.init rarr
rest = isolve k xs rxs
pairs = (find_pairs k x xs) + (find_pairs k rx rxs)
in
pairs + rest
solve :: Int -> [Int] -> Int
solve k xs =
let
arr = V.fromList $ sort xs
rarr = V.fromList $ sort xs
in isolve k arr rarr
main = do
[n, k] <- map read <$> words <$> getLine :: IO [Int]
xs <- map read <$> words <$> getLine :: IO [Int]
putStrLn $ show $ solve k xs