add solution week 270 task 2 in python

1 files changed, 51 insertions, 0 deletions

diff --git a/challenge-270/steven-wilson/python/ch-2.py b/challenge-270/steven-wilson/python/ch-2.py
new file mode 100644
index 0000000000..aa703848be
--- /dev/null
+++ b/challenge-270/steven-wilson/python/ch-2.py
@@ -0,0 +1,51 @@
+#!/usr/bin/env python3
+
+
+def distribute_elements(integers, x, y):
+    ''' Given an array of integers and two integers, x and y, execute one of
+    the two options:
+
+        Level 1:
+        Pick an index i of the given array and do integers[i] += 1
+
+        Level 2:
+        Pick two different indices i,j and do integers[i] +=1 and integers[j] += 1.
+
+    You are allowed to perform as many levels as you want to make every
+    elements in the given array equal. There is cost attach for each level,
+    for Level 1, the cost is $x and $y for Level 2.
+
+    In the end return the minimum cost to get the work done.
+
+    >>> distribute_elements([4, 1], x=3, y=2)
+    9
+    >>> distribute_elements([2, 3, 3, 3, 5], x=2, y=1)
+    6
+    '''
+    max_int = max(integers)
+    level_1_preferred = (x * 2 < y)
+    cost = 0
+    positions = [i for i, v in enumerate(integers) if v < max_int]
+    while positions:
+        if len(positions) > 1 and not level_1_preferred:
+            i, j = positions[:2]
+            integers[i] += 1
+            integers[j] += 1
+            cost += y
+            if integers[j] == max_int:
+                del positions[1]
+            if integers[i] == max_int:
+                del positions[0]
+        else:
+            i = positions[0]
+            integers[i] += 1
+            cost += x
+            if integers[i] == max_int:
+                del positions[0]
+    return cost
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(verbose=True)