-- SPDX-License-Identifier: Apache-2.0
-- Copyright 2025 ThreeSum Solution
--
-- Licensed under the Apache License, Version 2.0 (the "License");
-- you may not use this file except in compliance with the License.
-- You may obtain a copy of the License at
--
--     http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.

{-|
Module      : ThreeSum
Description : Solution to the 3Sum problem
Copyright   : (c) 2025
License     : Apache-2.0

This module provides an efficient solution to the 3Sum problem:
Given an integer array nums, return all the triplets [nums[i], nums[j], nums[k]]
such that i != j, i != k, and j != k, and nums[i] + nums[j] + nums[k] == 0.
The solution set must not contain duplicate triplets.
-}

module ThreeSum (threeSum, threeSumNaive, Triplet) where

import Data.List (sort, nub)
import Data.Set (Set)
import qualified Data.Set as Set

-- | Type alias for a triplet of integers
type Triplet = (Int, Int, Int)

-- | Efficient O(n²) solution using two pointers technique
-- This is the optimal approach for the 3Sum problem
threeSum :: [Int] -> [Triplet]
threeSum nums = removeDuplicates $ threeSum' sortedNums
  where
    sortedNums = sort nums
    
    threeSum' :: [Int] -> [Triplet]
    threeSum' [] = []
    threeSum' [_] = []
    threeSum' [_, _] = []
    threeSum' (x:xs) = findTriplets x xs ++ threeSum' (dropWhile (== x) xs)
    
    -- Find all triplets starting with x using two pointers
    findTriplets :: Int -> [Int] -> [Triplet]
    findTriplets x ys = twoSum x ys (reverse ys)
    
    -- Two pointers approach to find pairs that sum to -x
    twoSum :: Int -> [Int] -> [Int] -> [Triplet]
    twoSum _ [] _ = []
    twoSum _ _ [] = []
    twoSum x (l:ls) (r:rs)
      | l >= r = []  -- Pointers have met or crossed
      | sum' == target = (x, l, r) : twoSum x (dropWhile (== l) ls) (dropWhile (== r) rs)
      | sum' < target = twoSum x ls (r:rs)
      | otherwise = twoSum x (l:ls) rs
      where
        sum' = l + r
        target = -x

-- | Remove duplicate triplets from the result
removeDuplicates :: [Triplet] -> [Triplet]
removeDuplicates = Set.toList . Set.fromList . map normalizeTriplet
  where
    normalizeTriplet (a, b, c) = 
      let [x, y, z] = sort [a, b, c]
      in (x, y, z)

-- | Naive O(n³) solution for comparison and testing
-- This approach checks all possible triplets
threeSumNaive :: [Int] -> [Triplet]
threeSumNaive nums = removeDuplicates
  [ (nums !! i, nums !! j, nums !! k)
  | i <- [0..length nums - 3]
  , j <- [i+1..length nums - 2]
  , k <- [j+1..length nums - 1]
  , nums !! i + nums !! j + nums !! k == 0
  ]

-- | Example usage and test cases
main :: IO ()
main = do
  putStrLn "3Sum Problem Solutions in Haskell"
  putStrLn "=================================="
  
  -- Test case 1: Standard example
  let nums1 = [-1, 0, 1, 2, -1, -4]
  putStrLn $ "Input: " ++ show nums1
  putStrLn $ "Output (efficient): " ++ show (threeSum nums1)
  putStrLn $ "Output (naive): " ++ show (threeSumNaive nums1)
  putStrLn ""
  
  -- Test case 2: No solution
  let nums2 = [0, 1, 1]
  putStrLn $ "Input: " ++ show nums2
  putStrLn $ "Output: " ++ show (threeSum nums2)
  putStrLn ""
  
  -- Test case 3: All zeros
  let nums3 = [0, 0, 0]
  putStrLn $ "Input: " ++ show nums3
  putStrLn $ "Output: " ++ show (threeSum nums3)
  putStrLn ""
  
  -- Test case 4: Larger example
  let nums4 = [-2, 0, 1, 1, 2]
  putStrLn $ "Input: " ++ show nums4
  putStrLn $ "Output: " ++ show (threeSum nums4)