1. Analysis
(The code in the post was edited after I wrote this analysis, so the list of operations does not correspond to the latest version of the code. But the asymptotic result is the same.)
I'll write \$n\$ for the length of the score list, and \$k\$ for the number of team members.
On each iteration of the loop we have these operations:
- slices
score[m] and score[-m:];
- a list concatenation
score[m] + score[-m:];
- a call to
max on the concatenated list;
- a call to
score.index;
- a call to
score.pop.
Operations 1–3 take time proportional to \$m\$, and operations 4–5 take time proportional to \$n\$. The loop executes \$k\$ times, so the overall runtime is proportional to \$k(m + n)\$.
In the worst case, \$k\$ and \$m\$ are both \$Θ(n)\$, and so the overall runtime is \$Θ(n^2)\$, that is, quadratic in the size of the input.
We can demonstrate this behaviour by timing some random test cases:
from random import random
from timeit import timeit
def test(n, f=teamFormation):
score = [random() for _ in range(n)]
return timeit(lambda:f(score, n, n // 2 - 1), number=1)
and plotting the timings:
2. Better algorithm
In order to improve the performance of the code, we need to rewrite the code so that within the main loop, we avoid operations that take time proportional to the length of the slices, \$m\$, or to the length of the whole list, \$n\$.
What we'd like to do is to maintain a data structure containing the \$2m\$ values from the two ends of the list. The data structure must allow us to efficiently find and extract the maximum value. Then we can repeatedly extract the maximum value, and then replace it with the next value from the appropriate end of the remainder of list (the left end if the maximum value came from the left side, and right end otherwise).
The data structure we need for this is a heap, and Python provides us with the heapq module for manipulation of heaps.
I'll give the revised code below; read it and see if you can figure out how it works.
from heapq import heapify, heappop, heappush, nlargest
def best_team_score(score, k, m):
"""Return the maximum sum of k elements from the sequence score,
subject to the condition that only the first m and last m can be
selected at each stage.
"""
n = len(score)
if n < k:
raise ValueError("Not enough scores to make a team.")
if n <= 2 * m:
return sum(nlargest(k, score))
# Heap of (-score[i], i) pairs. The score is negated because
# Python's heaps are min-heaps, but we want the maximum score. By
# pairing with the index we can ensure that the leftmost score is
# returned in case of a tie, and that we can tell whether the
# maximum score came from the left or the right side.
heap = [(-score[j], j) for i in range(m) for j in (i, n - i - 1)]
heapify(heap)
# score[left:right] are the elements of score that have not yet
# been added to the heap.
left = m
right = n - m
total = 0
for _ in range(k):
# Leftmost maximum score in the heap.
s, i = heappop(heap)
total -= s
if right <= left:
# All scores are now in the heap.
pass
elif i < left:
# Maximum score came from left side.
heappush(heap, (-score[left], left))
left += 1
else:
# Maximum score came from right side.
assert right <= i
right -= 1
heappush(heap, (-score[right], right))
return total
We can check that this is correct by testing against the original code:
from itertools import product
from random import random
from unittest import TestCase
class TestBestTeamScore(TestCase):
def test_best_team_score(self):
for n in range(50):
score = [random() for _ in range(n)]
for k, m in product(range(1, n + 1), repeat=2):
found = best_team_score(score, k, k)
# teamFormation modifies score, so take a copy.
expected = teamFormation(score[:], k, k)
self.assertEqual(expected, found, score)
The heappush and heappop functions take time proportional to the logarithm of size of the heap, which is at most \$2m\$ so we expect the runtime to be propertional to \$k\log 2m\$, and since \$k\$ and \$m\$ are \$O(n)\$, the overall runtime is \$O(n\log n)\$:
3. Appendix: plotting
In case you're interested in how I drew the plots:
from functools import partial
import matplotlib.pyplot as plt
import numpy as np
def plot1():
emin, emax = 1, 5
x = 10 ** np.arange(emin, emax + 1)
y = np.vectorize(test)(x)
plt.xlabel("n")
plt.ylabel("t (seconds)")
a, = np.linalg.lstsq(x[:, np.newaxis] ** 2, y)[0]
xx = 10 ** np.arange(emin, emax, 0.01)
plt.loglog(xx, a * xx**2, label="least squares fit to $t = an^2$")
plt.loglog(x, y, 'r+', label="data")
plt.legend()
plt.grid(alpha=0.25)
plt.show()
def plot2():
emin, emax = 2, 7
x = 10 ** np.arange(emin, emax + 1)
y = np.vectorize(partial(test, f=best_team_score))(x)
plt.xlabel("n")
plt.ylabel("t (seconds)")
a, = np.linalg.lstsq((x * np.log(x))[:, np.newaxis], y)[0]
xx = 10 ** np.arange(emin, emax, 0.01)
plt.loglog(xx, a * xx * np.log(xx),
label="least squares fit to $t = an\log n$")
plt.loglog(x, y, 'r+', label="data")
plt.legend()
plt.grid(alpha=0.25)
plt.show()
