# List comprehension in Mathematica

Python supports a concept called "list comprehensions".

Following is sample example (written in Python syntax)

A = [[1, 2, 3], [2, 3, 4], [3, 4, 5]]
print( [ a + [i]  for a in A for i in range(1, 4) if i in a] )

[[1, 2, 3, 1], [1, 2, 3, 2], [1, 2, 3, 3], [2, 3, 4, 2], [2, 3, 4, 3], [3, 4, 5, 3]]


What's the most natural way to translate this Python code to Mathematica? I have the following method:

A = {{1, 2, 3}, {2, 3, 4}, {3, 4, 5}};
Select[Table[ a ~ Join ~ {i}, {a, A}, {i, 3}] ~ Flatten ~ 1,
MemberQ[ Most @ #, Last @ #]&]


but I think it's neither fast nor elegant, therefore I am looking for a better way.

• The same article you linked has a few lines about Mathematica too. – Dr. belisarius Nov 28 '13 at 12:05
• Maybe this? – Oleksandr R. Nov 28 '13 at 13:03

This one is very close to your Python code

Join @@ Table[Append[a, i], {a, A}, {i, Intersection[Range, a]}]

{{1, 2, 3, 1}, {1, 2, 3, 2}, {1, 2, 3, 3}, {2, 3, 4, 2}, {2, 3, 4, 3}, {3, 4, 5, 3}}

• After I overviewed this answer here, I acknowledged Mathematica's potential to build new styles -- instead of using array structures we could define {a,A} to be "for a in A" and "... if a in A" could be just the intersection. I think we could extend your answer here to create a visually-accessible Pythonic style to Mathematica. What do you think? – hhh Jun 27 '15 at 22:28
• This assumes that one can compute the set of valid a ahead of time, and then to implement some way to do that computation before performing the operation. A list comprehension is far more general. – Myridium Apr 28 at 0:53

Some of the other approaches might be much more efficient, but the following shows how one can create something which is probably as easy to read (if one is fluent in Mathematica) as a python list comprehension:

SetAttributes[listComprehend, HoldAll]

listComprehend[Verbatim[Condition][body_, crit_],iters:({_, __}..)] := Flatten[
Table[
If[crit, body, Unevaluated[Sequence[]]],
iterators
], 1]


use it like this:

listComprehend[Append[a, i] /; MemberQ[a, i], {a, A}, {i, 3}]


/; is the shortcut for Condition and usually can be read as "provided that" in mathematica code. Doing so, the above seems to be relatively clear code. Unfortunately the translation to one of the potentially more efficient approaches is not that simple...

A = {{1, 2, 3}, {2, 3, 4}, {3, 4, 5}};
Join @@ Table[If[a~MemberQ~i, a~Join~{i}, Unevaluated[]], {a, A}, {i, 3}]
Join @@ Table[a~Join~{i}, {a, A}, {i, Select[Range@3, a~MemberQ~# &]}]


{{1, 2, 3, 1}, {1, 2, 3, 2}, {1, 2, 3, 3}, {2, 3, 4, 2}, {2, 3, 4, 3}, {3, 4, 5, 3}}

• +1 for writing Unevaluated[] instead of Unevaluated[Sequence[]], or ##&[] :) – Jacob Akkerboom Nov 28 '13 at 14:32
• @Jacob What's wrong with my "vanishing function" ##&[]? (+1, by the way) – Mr.Wizard Nov 28 '13 at 16:30
• @Mr.Wizard haha nothing really, in fact I was glad to show off that I knew about it :P. But also despite all the talk about this, I may have missed the option of just using Unevaluated[], so I guess I learned something :). – Jacob Akkerboom Nov 28 '13 at 16:49
• @Jacob I attempted to cover that in this answer but I guess I should have been more explicit. I said that "All of these heads meet this requirement:" and listed Unevaluated and ## & -- both would be used the same way: Unevaluated[] and ## &[]. – Mr.Wizard Nov 29 '13 at 12:20
• @Mr.Wizard that one is a nice answer (as I suppose you can tell by the badge :P). I'm not very focussed right now, but perhaps an additional point of confusion could by that neither Apply[] nor Symbol[] evaluate to Sequence[] (anymore?). – Jacob Akkerboom Nov 29 '13 at 12:42

Not really sure if you're after the functionality, the way to generate your list, or some performance improvement. Anyway:

Cases[Outer[Join, A, List /@ Range@Length @A, 1], {___, x_, ___, x_}, 2]
(*
{{1, 2, 3, 1}, {1, 2, 3, 2}, {1, 2, 3, 3},
{2, 3, 4, 2}, {2, 3, 4, 3}, {3, 4, 5, 3}}
*)


Edit

For efficiency, you may want to generate only the final elements, without the need to post filter them:

Flatten[Function[{u}, Table[Join[u, {i}], {i, u}]][#] & /@ A, 1]


or:

Flatten[Function[{u}, ReplaceList[#, Thread[x_ :> Flatten[{x, #}] & /@ u]]][#] & /@ A, 1]


If you are going for a literal translation from Python to Mathematica, I think that something like the following snippet should do the trick:

Part[
Reap@Do[
Do[
If[
MemberQ[a, r],
Sow@Flatten[{a, r}]
], {r, R}], {a, A}], 2, 1]


where A = {{1, 2, 3}, {2, 3, 4}, {3, 4, 5}} and R=Range[1,3] (note how in Python $list(range(1, 4))=[1, 2, 3]$)

Nevertheless, this question is about finding "...the most natural way to translate this Python code to Mathematica" and my view is that going for the literal translation is not always the best option (although it could very possibly be what you need at a certain point in time).

The answers to this question provide useful examples of how versatile Mathematica is.

Another natural way to achieve the intended result, is to use Map:

Flatten[
Map[(
a = #;
Flatten[{a, #}] & /@ Pick[R, MemberQ[a, #] & /@ R]
) &, A], 1]


Perhaps, a more concise way to achieve the same result is the following:

f[a_, r_] := If[MemberQ[a, r], Flatten[{a, r}]]

Cases[
Flatten[
Outer[f, A, R, 1], 1], Except[Null], 1]


I think that this way is more close in practice to what the list comprehension does: it traverses the items of A (rows) and then it traverses the items in the range until it finds a match.Then it appends the range item to the list row and proceeds to the next row.