PepeToro
PepeToro

Reputation: 559

solve particular symbolic equation in python

I want to solve the following problem with python, if possible with sympy.

Let n be a fixed positive number. Let p=(p_1,...p_n) be a fixed known vector of positive integers. Let d be a fixed, known positive integer. Let q=(q_1,...,q_n) be a vector of unknown nonnegative integers.

How can I get all the solutions of p.q=d?

Where . means dot product.

Actually I can solve this for each individual n. But I want to create a function

def F(n,p,d):
...
return result

Such that result is a, e.g., list of all solutions. Note that from the restrictions made above, there is a finite number of solutions for each triplet of data (n,p,d).

I can't figure a way to do this, so any suggestion will be appreciated.


Added.

Example: suppose n=3 (the case n=2 is trivial), p=(2,1,3), d=3. Then I would do something like

res=[]
for i in range (d):
    for j in range (d):
        k=d-p[0]*i-p[2]*j
        if k>=0:
            res.append([i,k,j])

Then res=[[0, 3, 0], [0, 0, 1], [1, 1, 0]] which is correct.

As you can imagine, the bigger n is, the more for loops I need if I want to follow the same idea. So I do not think this is a good way to do it for arbitrary n, say n=57 or whatever big enough...

Upvotes: 1

Views: 158

Answers (1)

Joel Cornett
Joel Cornett

Reputation: 24788

Following the algorithm you provided:

from itertools import product
dot = lambda X, Y: sum(x * y for x, y in zip(X, Y))
p = [1, 2, 3, ...] # Whatever fixed value you have for `p`
d = 100 # Fixed d
results = []
for q in product(range(0, d+1), repeat=len(p)):
    if dot(p, q) == d:
        results.append(q)

However this is slightly inefficient since it is possible to determine prior to computing the entire dot product, whether k will be positive. So let's define the dot product like this:

def dot(X, Y, d):
    total = 0
    for x, y in zip(X, Y):
        total += x * y
        if total > d:
            return -1
    return total

Now, as soon as the total exceeds d, the calculation exits. You can also express this as a list comprehension:

results = [q for q in product(range(0, d+1), repeat=len(p)) if dot(p, q, d) == d]

Upvotes: 3

Related Questions