Input list has n elements, each element has a score. The goal is to assemble a group (of size ‘team’) of elements together and to find their total score.

Inputs:

  • score [] of size n,
  • team which is the team size,
  • m tells us how to slice the list

Selection of elements from the list is done in the following way:

  1. During each selection, we have to choose an element with the highest score among either the first m in available elements or the last m available elements in the list. The element is then removed from the list and added to the team total score.
  2. If there are multiple elements with the same highest score among the first or last in available elements, the element closest to the start of the line is picked.
  3. If there are less than in available elements. then we pick the highest scored element from all available elements.

Sample Input and Output:

Sample 1

n:9, score:[17,12,10,2,7,2,11,20,8,],team:3,m:4
OUTPUT : 49

Sample 2

n:8,score:[18,5,15,18,11,15,9,7],team:5,m:1
OUTPUT : 60

Need an optimal solution for this Python code. I am getting the desired output. But I am looking for a better solution.

def teamFormation(score, team, m):
    total_score = 0
    for _ in range(team):
        if len(score) < 2*m:
            total_score += score.pop(score.index(max(score)))
        elif max(score[:m]) > max(score[-m:]):
            total_score += score.pop(score[:m].index(max(score[:m])))
        else:
            total_score += score.pop(score[-m:].index(max(score[-m:]))-m)
    return total_score
  • 5
    Welcome to Code Review. While it's fine to update your code as long as there hasn't been a review yet, keep in mind that you're invalidating possible in-progress reviews. After you get answers, do not update the code in your question to incorporate feedback from answers. – Zeta Feb 27 at 11:12
  • 1
    Please don't update the code in the question after your code has been reviewed. – Peilonrayz Feb 28 at 11:02
  • Sorry mate! It was just a small renaming. Didn't touch the logic though. – Johnny Feb 28 at 11:04
up vote 10 down vote accepted

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:

  1. slices score[m] and score[-m:];
  2. a list concatenation score[m] + score[-m:];
  3. a call to max on the concatenated list;
  4. a call to score.index;
  5. 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:

Graph showing quadratic behaviour of OP

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)\$:

Graph showing n log n behaviour of revised code

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()
  • @Peilonrayz: There are two points to the assertion. First, it checks that I programmed it right (since there's an opportunity to make an off-by-one error somewhere). Second, it answers the question that a reader might ask: "but what if left <= i < right?" by asserting that this never happens. – Gareth Rees Feb 27 at 13:26
  • Thank you, those reasons make sense. I'll have to start doing things like that at some point too. – Peilonrayz Feb 27 at 13:34
  • 1
    Thank you very much for the detailed explanation and the solution! Helped a lot! I am extremely grateful. – Johnny Feb 27 at 16:42

To improve your current code, you can:

  • Use itertools to build an iterable to get the index and value from ms in one go.
  • You don't need your branching, as it duplicates code.
  • You can use yield to remove the sum logic from the code.

Getting:

from itertools import chain, count
from operator import itemgetter


def team_formation(score, team, m):
    value_getter = itemgetter(1)
    for _ in range(team):
        ms = chain(
            zip(count(), score[:m]),
            zip(count(len(score-m)), score[-m:])
        )
        index, value = max(ms, key=value_getter)
        yield value
        score.pop(index)

def teamFormation(score, team, m):
    return sum(team_formation(score, team, m))

This however isn't that great. It's \$O(t(m + n))\$ time complexity. You can remove the \$O(tn)\$ from this by removing the need for score.pop. The simplest way to do this is to just keep ms built, and make score a deque.

import collections
from operator import itemgetter

def team_formation(score, team, m):
    value_getter = itemgetter(1, 0)
    score = collections.deque(score)
    fns = [collections.deque.popleft, collections.deque.pop]
    ms = []
    try:
        for place in [True, False]:
            for _ in range(m):
                ms.append((fns[place](score), place))
    except IndexError:
        pass

    for _ in range(team):
        index, (value, place) = max(enumerate(ms), key=value_getter)
        ms.pop(index)
        yield value
        if score:
            ms.append((fns[place](score), place))

This gives \$O(tm + m + n)\$ runtime. You can possibly further improve performance by using another datatype for ms. I was going to use the standard library heap, but as it's not stable I didn't. Using it would allow you to change ms to get \$O(t\log{m} + m + n)\$. Which would be a drastic increase in performance. There however is actually a way to get around this, as pointed out in Gareth Rees answer.

  • The first version would benefit from mentioning itertools.izip; and maybe looking at itertools.islice as well. The second would benefit from using a generator comprehension instead of the nested for loops, or factoring that out into a function and have it yield values. I also have always found indexing using booleans to be hacky at best in Python; making a generator function would also allow skipping that step. Lastly, fns should be a tuple instead of a list. – Dannnno Feb 27 at 12:43
  • @Dannnno 1. I assume Python 3 so izip is zip. 2. How would islice help? Are you saying to get \$O(n)\$ rather than \$O(m)\$ complexity? 3. I'm not sure what you're talking about. 4. bool is a subclass of int, I don't see what's hacky. 5. fns can become a tuple if you want, but meh – Peilonrayz Feb 27 at 12:47

Removing duplicated logic

By using temporary variables, you could make your code easier to read and avoid the amount of duplicated logic. Also, that should gives a tiny performance boost:

def teamFormation(score, team, m):
    total_score = 0
    for _ in range(team):
        if len(score) < 2*m:
            total_score += score.pop(score.index(max(score)))
        else:
            beg, end = score[:m], score[-m:]
            max_beg, max_end = max(beg), max(end)
            if max_beg > max_end:
                index_max = beg.index(max_beg)
            else:
                index_max = end.index(max_end) - m
            total_score += score.pop(index_max)
    return total_score
  • You're most welcome. My answer is actually pretty poor compared to the others but I hope it will be useful for you :) – Josay Feb 27 at 16:48

Your Answer

 

By clicking "Post Your Answer", you acknowledge that you have read our updated terms of service, privacy policy and cookie policy, and that your continued use of the website is subject to these policies.

Not the answer you're looking for? Browse other questions tagged or ask your own question.