Search code examples
pythonarraysoptimizationnested-loops

Optimization of Python Nested loops with Two Arrays


In an assessment, I got the following question:

"There are two teams of size N, and the skill values of each player on the team are stored inside an array, write an algorithm that counts all the times that the number one team wins a round. To find out if team number one has won a round, it is necessary to compare pairs of players between each of the two teams, taking into account that the comparisons can be made in the form (0<=i<j<n). In essence, if and only if group1[i]+group1[j ] > group2[i]+group2[j] then team one wins"

I came out with a solution of O(n^2) time complexity which is the following:

def countWins(group1, group2):
    n = len(group1)
    wins = 0
    for i in range(n):
        for j in range(i+1, n):
            if group1[i] + group1[j] > group2[i] + group2[j]:
                wins += 1
    return wins

Despite the simplicity of the algorithm, I will explain it briefly:

The first element of the pair is collected through the outermost loop, while thanks to the innermost one the second element of the pair is collected without considering those in a position preceding or equal to "i".

For each detected pair the win counter is incremented if it satisfies the track condition.

What I thought of to decrease the complexity of the algorithm was to use another data structure, but I would still have to go through the two arrays, which would have made using any structure other than the one I have inevitably less efficient (I might be wrong).

I also thought that sorting the arrays might help arrive at a better solution, but nothing useful came to mind.


Solution

  • let's change the problem, instead of finding the locations where

    group1[i] + group1[j] > group2[i] + group2[j]
    

    we can move the right side over to the left side

    (group1[i] - group2[i]) + (group1[j] - group2[j]) >  0
    

    then if we just subtract both arrays

    differences = group1 - group2 # assume this works
    differences[i] + differences[j] > 0
    

    now the problem is to find the number of pairs (differences[i], differences[j]) which sum greater than 0, this problem is well known, i might just borrow its solution from geek for geeks number of pairs greater than 0 and convert it to python as follows:

    import bisect
    
    def countWins(group1, group2):
        wins = 0
        differences = [x-y for x,y in zip(group1,group2)]
        differences.sort()
        a = differences
        n = len(differences)
        # Loop to iterate through the array
        for i in range(n): 
     
            # Ignore if the value is negative
            if (a[i] <= 0):
                continue
     
            # Finding the index using lower_bound
            j = bisect.bisect_left(a, -a[i] + 1);
     
            # Finding the number of pairs between
            # two indices i and j
            wins += i - j;
        return wins
    
    def main():
        arr1 = [1,3,4,6]
        arr2 = [0,1,4,7]
        print(countWins(arr1,arr2))
    
    main()
    
    4
    

    note that the algorithm for finding the number of pairs greater than 0 could be further optimized, but this solution is already O(n*logn), further optimizations won't change that, it will only chop off a few instructions.