haskelllist-comprehension

Returning a variable only after a successful test in a list comprehension


I currently have the following code (which works):

primL :: [Int]
primL = [2*(m^2+n^2) + 2*b|
    m<-[2..12909],
    n<-[1..m-1],
    m>n,
    odd (m+n),
    gcd m n == 1,
    let b | 4*m*n == m^2+n^2+1 || 4*n*m == m^2+n^2-1 = 2*m*n
          | 2*(m^2-n^2) == m^2+n^2+1 || 2*(m^2-n^2) == m^2+n^2-1 = m^2-n^2
          | otherwise = 0,
    b /= 0]

Is there a way to avoid that last bit of throwing a fake value for b? I would like to be able to do the test, if it fails then the next values for m and n are tried, and if it succeeds, then b is defined accordingly.

The only alternative I have found so far is to do the tests twice, once as part of the list definition, and a second time below to define b if the test succeeded, but I think that's not great performance-wise.


Solution

  • I wouldn't worry about repeating calculations too much, if that feels like the most natural way to write this code. List allocation costs will trump additions and multiplications like crazy. (In fact, even your existing calculations duplicate some work, as there are several shared subterms.)

    If you really want to try to avoid repeated calculations but keep the approach you proposed, it's not hard -- just name the calculations you'll do multiple times. So, you could write, for example:

    let prod    = 2*m*n
        sumSq   = m^2 + n^2
        diffSq  = m^2 - n^2
        branch0 = abs (2*prod   - sumSq) == 1
        branch1 = abs (2*diffSq - sumSq) == 1
        b = if branch0 then prod else diffSq,
    branch0 || branch1]
    

    With GHC, each variable's value will be calculated at most once; sometimes zero times (e.g. if firstBranch is true, then diffSq will not be calculated). As a bonus, you can include sumSq in your list comprehension head to remove a recalculation there, too.