Python: Writing quite complicated piece of code as a List Comprehension

Question

I have solved the Word Break problem with the following code:

def wordBreak(self, s: str, wordDict: List[str]) -> bool:
        n = len(s)
        dp = [False] * n
        for i in range(n):
            for j in range(i + 1, n + 1):
                if (s[i : j] in wordDict) and (i == 0 or dp[i - 1]):
                    dp[j - 1] = True
        return dp[n - 1]

And I'm trying to write it as a list comprehension to make it "more or less one-line code".

I have been trying to express dp as follows:

dp = [True if (s[i:j] in wordDict) and (i==0 or dp[i-1]) for j in range(i+1, n+1) for i in range(n) else False]

However, I kind of got stuck because in the original code it is set dp[j - 1] = True while I couldn't make it with a list comprehension.

I know it's not a good idea to write a big piece of code as a list comprehension (I'm confused already too) but it is just for educational purposes.

Answer 1

Indeed, as dp is written to, and not only that, but those freshly written values are read again in next iterations, this would not be a candidate for list comprehension. If however you are willing to sacrifice some best practices in Python, you can get close.

First, you can turn the inner loop into a list comprehension:

def wordBreak(self, s: str, wordDict: List[str]) -> bool:
    n = len(s)
    dp = [False] * n
    for i in range(n):
        dp[i:] = [dp[j] or
                   (s[i:j+1] in wordDict) and (i == 0 or dp[i-1])
                   for j in range(i, n)]
    return dp[-1]

Let's slightly change the algorithm, so the list overwriting happens from index 0, and not from index i. Instead we shorten the list gradually. At the same time we designate dp[0] as the value to read from the previous iteration, so we also prefix it to the initial list:

def wordBreak(self, s: str, wordDict: List[str]) -> bool:
    n = len(s)
    dp = [True] + [False] * n
    for i in range(n):
        dp = [dp[j-i+1] or s[i:j+1] in wordDict and dp[0]
                   for j in range(i, n)]
    return dp[0]

But dp = is an assignment that in principle blocks a wider use of list comprehension. You can get around this with a a-Pythonic function, which both alters and returns something:

def assign(self, target, source):
    target[:] = source  # mutate the target
    return target[-1]  # for our purposes we only need it to return the last value

And now we can include it in a larger list comprehension expression:

def wordBreak(self, s: str, wordDict: List[str]) -> bool:
    n = len(s)
    dp = [True] + [False] * n
    return [
        self.assign(dp, [dp[j-i+1] or s[i:j+1] in wordDict and dp[0]
                        for j in range(i, n)])
                        for i in range(n)
    ][-1]

But to repeat: this is not pythonic. It is best practice to let a function either mutate an argument (or self), or return a value, but not both. There are only a few exceptions to this rule (e.g. .pop())

If instead of list comprehension, you are happy with a functional expression, an alternative could be to use functools.reduce:

def wordBreak(self, s: str, wordDict: List[str]) -> bool:
    n = len(s)
    return reduce(lambda dp, i: [dp[j-i+1] or s[i:j+1] in wordDict and dp[0]
                                 for j in range(i, n)], 
                  range(n), 
                  [True] + [False] * n)[0]