haskell - function with lists involving Maybe not working

Question

I have the following function:

-- xs: list to be changed
-- ws: list of indices where the values will change to 0
replaceAtIndices xs ws = [if (fromJust (elemIndex x xs)) `elem` ws then 0 else x | x <- xs]

The function takes in 2 lists. ws are the indices of the values in xs that I want to change to 0.

For some reason, it works for some cases, and not for others:

*Main> replaceAtIndices [1,2,3,4] [2,3]

[1,2,0,0] -- correct

*Main> replaceAtIndices [1,2,3,4] [2]

[1,2,0,4] -- correct

*Main> replaceAtIndices [1,1,2,1,3] [3]

[1,1,2,1,3] -- SHOULD be [1,1,2,0,3]

Can anyone please explain why this is?

Thanks in advance!

Answer 1

elemIndex returns the index of the first occurrence of the item in the list, so in the third case it always returns 0 for the index of 1 which doesn't match 3 so nothing gets replaced.

A better way of associating indexes with the items is to use zip :

replaceAtIndices xs ws = [if i `elem` ws then 0 else x | (i, x) <- zip [0..] xs]