CODEWITHSUNDEEP
haskell
Norrius
Norrius

Reputation: 7930

Zipping two HashMaps

Let's say I have two HashMaps:

import Data.HashMap.Strict
m1 :: HashMap Char Int
m2 :: HashMap Char Int
??? :: (v1 -> v2 -> v3) -> HashMap k v1 -> HashMap k v2 -> HashMap k v3
-- or even
??? :: (k -> v1 -> v2 -> v3) -> HashMap k v1 -> HashMap k v2 -> HashMap k v3

I'd like some kind of zip to walk over these maps, where:

  1. if a key exists in one map but not the other, it is substituted with a default (sort of "zip longest")
  2. the result type can be different from the original types.

This matches the signature of intersectionWithKey with semantics closer to unionWithKey.

A concrete example: for a simple case of applying (,) with a 0 default, I can do

unionWith add (map (,0) m1) (map (0,) m2) :: HashMap Char (Int, Int)
add (a, x) (b, y) = (a+b, x+y)

This provokes the question: is there a clean way to do this in general?

Upvotes: 2

Views: 290

Answers (2)

chepner
chepner

Reputation: 531958

In general:

  1. Merge the two maps to get a list of all keys
  2. Define a new function that, given a key, uses lookupDefault to get the appropriate value (default or real) from each map in order to compute the value for the new map
  3. Map that function over the keys, and create a new hash map from the result.
zipMap :: (v1 -> v2 -> v3) -- Combining function
       -> v1               -- default value for keys unique to the second map
       -> v2               -- default value for keys unique to the first map
       -> HashMap k v1
       -> HashMap k v2
       -> HashMap k v3
zipMap f d1 d2 m1 m2 = let allKeys = keys (union m1 m2)
                           f' k = (k, f (lookupDefault d1 k m1) (lookupDefault d2 k m2))
                       in fromList $ map f' allKeys

Upvotes: 1

danidiaz
danidiaz

Reputation: 27766

The these package provides the These datatype, which is a bit like Either but has an extra constructor for when both cases are present.

The package also defines typeclasses for "zipping with padding" over data structures. In particular the Align typeclass, which has a HashMap instance.

For example, to zip two maps using 0s for missing entries:

import Data.These
import Data.Align

padWithZeros :: (Align f, Num a, Num b) => f a -> f b -> f (a, b)
padWithZeros = alignWith (fromThese 0 0)

There are also Data.Align.Indexed and Data.Align.Keyed, which give the zipping function access to the key.

One disadvantage of these is that the dependency foot print is somewhat heavy.

Upvotes: 1

Related Questions