Update maximum-cost-of-trip-with-k-highways.py

Author: kamyu104

Python/maximum-cost-of-trip-with-k-highways.py

@@ -5,7 +5,7 @@
 
 
 # combination based dp
-class Solution2(object):
+class Solution(object):
     def maximumCost(self, n, highways, k):
         """
         :type n: int
@@ -19,24 +19,25 @@ def maximumCost(self, n, highways, k):
         for c1, c2, t in highways:
             adj[c1].append((c2, t))
             adj[c2].append((c1, t))
+        result = -1 if k != 1 else 0
         dp = [[0, []] for _ in xrange((1<<n))]
         for i in xrange(n):
             dp[1<<i][1].append(i)
-        result = -1
         for cnt in xrange(1, n+1):
             for choice in itertools.combinations(xrange(n), cnt):
                 mask = reduce(lambda x, y:x|(1<<y), choice, 0)
-                for u in dp[mask][1]:
+                total, lasts = dp[mask]
+                for u in lasts:
                     for v, t in adj[u]:
                         if mask&(1<<v):
                             continue
                         new_mask = mask|(1<<v)
-                        if dp[mask][0]+t < dp[new_mask][0]:
+                        if total+t < dp[new_mask][0]:
                             continue
-                        if dp[mask][0]+t == dp[new_mask][0]:
+                        if total+t == dp[new_mask][0]:
                             dp[new_mask][1].append(v)
                             continue
-                        dp[new_mask][0] = dp[mask][0]+t
+                        dp[new_mask][0] = total+t
                         dp[new_mask][1] = [v]
                         if bin(mask).count('1') == k:
                             result = max(result, dp[new_mask][0])