Restoring addresses with Haskell

Leetcode has a very nice problem called Restore IP Addresses. The task is that for given string s containing only digits, we find all possible valid IPv4 addresses that can be formed by inserting dots into s. For example, from "25525511135" we can form only addresses 255.255.11.135 and 255.255.111.35.

Solution in 16 lines

First, we will construct three functions of type String -> [(Int, String)]. These functions from the start of a given string try to parse a single-digit, a two-digit, or a three-digit number. If parsing is successful, a singleton list is returned. The only element of the list is a pair of the parsed number and the rest of the string. If parsing fails, an empty list is returned.

parse1, parse2, parse3 :: String -> [(Int, String)]

parse1 (x:xs) = [(read [x], xs)]
parse1 _ = []

parse2 (x:y:xs) = [(read [x, y], xs) | x /= '0']
parse2 _ = []

parse3 (x:y:z:xs) = [(num, xs) | let num = read [x, y, z], num <= 255, x /= '0'] 
parse3 _ = []

The reason we are returning a list of pairs instead a pair is that we want somehow to signal parsing failure. In some sense, it would be more correct to return a pair wrapped in Maybe, but I wanted to keep things simple.

Function parse1 can only fail if an empty string is passed. Otherwise, it will parse a single digit. For example parse1 "12345" gives us [(1, "2345")]. Similarly, parse2 will fail if a passed string contains less than two characters. But it will also fail if the first character of a string is '0'. Function parse3 will fail if the provided string doesn’t contain at least 3 characters, if the provided string starts with a '0', or if the parsed number is bigger than 255.

We will apply functions parse1, parse2 and parse3 to some string s and get three lists. Concatenating results we are left with a single list of successful parses. We can use concatMap = concat . map to write things more concisely:

successes s = concatMap ($ s) [parse1, parse2, parse3]

Note that here we are using the application operator to apply value to each function in list. Construction map ($ x) [f1, f2 .. fn] is same as [f1 x, f2 x .. fn x].

For example, successes "3211" will give us [(3,"211"), (32,"11")]. Parser parse3 failed in this case because 321 > 255, so we have only two results in the returned list.

The only thing left is to recursively restore the rest of the string. For that, we will construct a function restore :: String -> [[Int]], which gives a list of restored addresses. Restored addresses will be given as values of type [Int].

If (n, rest) is a successful parse of a string s, then we can continue restoring the string rest. This will give us a list l of restored addresses for string rest, so restored addresses of strings s that start with octet n are constructed by prepending n to each element of l. In the end, we are left with an empty string. In that case, we return a list that contains an empty list. This empty list is a “seed” on which we prepend successfully parsed integers as we go back through recursive calls.

restore :: String -> [[Int]]
restore "" = [[]]
restore s  = concatMap (\(n, rest) -> map (n :) $ restore rest) successes
    where successes =  concatMap ($ s) [parse1, parse2, parse3]

Returning to given example, successes is [(3,"211"), (32,"11")]. Recursively restoring "211" should give us [[21,1], [2,1,1], [2,11]]. If we prepend 3 at each element we will get [[3,21,1], [3,2,1,1], [3,2,11]]. Doing the same for (32,"11") gives us [[32,1,1], [32,11]]. Finally, we just concatenate these lists.

If we get to the end, for example if (n, rest) is (1,""), then recursive call restore "" will give us [[]]. Then the map will add 1 to each element of the returned list, which will give us [[1]]. That is exactly what we want! If we returned an empty list [], then map (n :) $ restore rest would evaluate to the [].

Let’s test our function with the first provided test case:

ghci> restore "25525511135"
[[2,5,5,2,5,5,1,1,1,3,5],[2,5,5,2,5,5,1,1,1,35],[2,5,5,2,5,5,1,1,13,5], ...

Oh, snap! We got a list of hundreds results, but we expected only two. The problem is that program doesn’t have a limit on the number of octets. In order to solve this, we want to add an additional parameter to the restore function. This parameter signifies how much more octets we want to parse.

If we don’t want to parse any more octets and we have an empty string to restore, then everything is OK, and we return our seed [[]].
If we don’t want any more octets to parse but we have a nonempty string, we parsed “too little” earlier in the stack. Similarly, if have nothing more to parse, but we expect more octets, we parsed “too much” earlier. In both cases we return [].
In all other cases, we do what we did: calling restore recursively, but this time reducing the first parameter of the call each time.

restore :: Int -> String -> [[Int]]
restore 0 "" = [[]]
restore 0 _  = []
restore _ "" = []
restore d s  = concatMap (\(n, rest) -> map (n :) $ restore (d - 1) rest) successes
    where successes =  concatMap ($ s) [parse1, parse2, parse3]

Now function works perfectly:

ghci> restore 4 "25525511135"
[[255,255,11,135],[255,255,111,35]]
ghci> restore 4 "101023"
[[1,0,10,23],[1,0,102,3],[10,1,0,23],[10,10,2,3],[101,0,2,3]]