NOTE: This article is in draft status as part of my Write in Public initiative. Expect major changes to the content and format in the future.
Introduction
Problem Statement
Let’s say that number
a
feels comfortable with numberb
ifa ≠ b
andb
lies in the segment[a - s(a), a + s(a)]
, wheres(x)
is the sum ofx
’s digits.
How many pairs(a, b)
are there, such thata < b
, botha
andb
lie on the segment[l, r]
, and each number feels comfortable with the other (soa
feels comfortable withb
andb
feels comfortable witha
)?Example
Forl = 10
andr = 12
, the output should be
solution(l, r) = 2
.
Here are all values ofs(x)
to consider:
•s(10) = 1
, so10
is comfortable with9
and11
;
•s(11) = 2
, so11
is comfortable with9
,10
,12
and13
;
•s(12) = 3
, so12
is comfortable with9
,10
,11
,13
,14
and15
.
Thus, there are2
pairs of numbers comfortable with each other within the segment[10; 12]
:(10, 11)
and(11, 12)
.
Input/Output
• [execution time limit] 4 seconds (py3)
• [memory limit] 1 GB
• [input] integer l
Guaranteed constraints:
1 ≤ l ≤ r ≤ 1000
.
• [input] integer r
Guaranteed constraints:
1 ≤ l ≤ r ≤ 1000
.
• [output] integer
The number of pairs satisfying all the above conditions.
My Solution
Initial Attempt
def solution(l, r):
pairs = []
# loop over all possible a's
for a in range(l, r+1):
# calculate the sum of digits
for b in range(l, r+1):
if a != b and b in range(a-s(a),a+s(a)+1):
pairs.append((a,b,))
return len(list(filter(lambda x: (x[1],x[0]) in pairs, pairs)))//2
def s(x):
return sum(map(lambda x: int(x), str(x)))
This solution will work if time is not a concern. This has $O(n^2)$ time complexity which is too much to pass the test.
Round 2: List Comprehension
In general, list comprehensions tend to perform faster than nested loops. ^[Citation needed]
def solution(l, r):
pairs = [(a,b,) for a in range(l, r+1) for b in range(l,r+1) if (a!=b and b in range(a-s(a),a+s(a)+1))]
return len(list(filter(lambda x: (x[1],x[0]) in pairs, pairs)))//2
def s(x):
return sum(map(lambda x: int(x), str(x)))
But this doesn’t pass any additional tests that didn’t pass with the nested loop approach.
Round 3: Precalculating…
I can save a little processing time by calculating the segment before the inner loop…
def solution(l, r):
pairs = []
for a in range(l, r+1):
segment = list(range(a - s(a), a + s(a) + 1))
for b in range(l, r+1):
if (a != b) and (b in segment):
pairs.append((a,b,))
return len(list(filter(lambda x: (x[1],x[0]) in pairs, pairs)))//2
And that does pass Test 5. But other tests are not passed.
Round 4: Move
After rereading the problem statement, I saw the condition a < b
. This made me think that I could reduce the search space by swapping the inner loop with the outer loop.
Thus, the solution that passes the sample tests…
def solution(l, r):
count = 0
for b in range(l+1, r+1):
for a in range(l,b):
a_segment = range(a - s(a), a + s(a) + 1)
b_segment = range(b - s(b), b + s(b) + 1)
if (b in a_segment) and (a in b_segment):
count += 1
return count
def s(x):
return sum(map(lambda x: int(x), str(x)))
Using the DRY principle, and optimizing for use of list comprehensions…
def solution(l, r):
return sum([1 for b in range(l+1, r+1) for a in range(l,b)
if (a in segment(b)) and (b in segment(a))])
def segment(x):
sum_of_digits = sum(map(lambda x: int(x), str(x)))
return range(x - sum_of_digits, x + sum_of_digits + 1)
This solution sacrifices a little bit of readability in exchange for an increase in performance
Round 5: The Gempire Strikes Back
Now here’s a trick: Memoization
By caching previously calculated values, we can gain some speed enhancements on the s
function.
def solution(l, r):
count = 0
for b in range(l+1, r+1):
for a in range(l,b):
a_segment = range(a - s(a), a + s(a) + 1)
b_segment = range(b - s(b), b + s(b) + 1)
if (b in a_segment) and (a in b_segment):
count += 1
return count
def memoize(func):
memoized = {}
def inner(number):
if number not in memoized:
memoized[number] = func(number)
return memoized[number]
return inner
@memoize
def s(x):
return sum(map(lambda x: int(x), str(x)))
And we pass all tests!
Further Optimization
However, there may still some space inefficiency here. Instead of using the range
function, let’s use a tuple indicating the bounds of the segment.
def solution(l, r):
count = 0
for b in range(l+1, r+1):
for a in range(l,b):
if is_in_segment(b,calculate_segment(a)) and is_in_segment(a,calculate_segment(b)):
count += 1
return count
def memoize(func):
memoized = {}
def inner(number):
if number not in memoized:
memoized[number] = func(number)
return memoized[number]
return inner
@memoize
def calculate_segment(x):
sum_of_digits = sum(map(lambda x: int(x), str(x)))
return (x - sum_of_digits, x + sum_of_digits)
def is_in_segment(number, segment):
return (number >= segment[0]) and (number <= segment[1])
[!Idea] Range membership test may be faster than boundary check https://realpython.com/python-range-membership-test/
Solution from nixerasse using the itertools
module
import itertools
def solution(L, R):
def is_comfortable(a,b):
s = sum(int(d) for d in str(a))
return a - s <= b <= a + s
cnt = 0
for a, b in itertools.combinations(range(L, R + 1), 2):
if is_comfortable(a, b) and is_comfortable(b, a):
cnt += 1
return cnt