Search code examples
pythonlistsum

Sum based on a criterion


I have a data sheet where I want to do a calculation until I find a zero in one specific column. Then I want to sum all the results of this calculation up to that zero and save the result in an array. I tried with an np array and a list but it does not work:

tst = []
x = data[1:len(data),0]
y = data[1:len(data),1]  
intt = data[1:len(data),2]
for i in range(0,len(data)):
   if intt[i]!=0:         
      tst.append((x[i]**2.0+ y[i]-y[i-1])**2.0)

I want the tst to contain in each position, the sum of the expression in the append().

Thanks!


Solution

  • I looks like you skip the first datapoint since you write x = data[1:len(data),0] and so forth. As @trincot mentioned, you also have to care about the y[i-1] case for i=0. Maybe the following will help you:

    tst = []
    x = data[:,0]
    y = data[:,1]  
    intt = data[:,2]
    total = []
    for i in range(1,len(data)):
       if intt[i]!=0:         
          tst.append((x[i]**2.0+ y[i]-y[i-1])**2.0)
       else:
          total.append(sum(tst))
          break
    total = np.array(total)
    

    This includes all data points in x, y, and intt, but the first data point will still be skipped since the loop starts with i=1.

    Edit: A cleaner solution, based on comment from @Barmar.

    tst = []
    x = data[:,0]
    y = data[:,1]  
    intt = data[:,2]
    for i in range(1,len(data)):
       if intt[i]!=0:         
          tst.append((x[i]**2.0+ y[i]-y[i-1])**2.0)
       else:
          total = np.array(sum(tst))
          break
    

    Edit: Another solution that goes through all of intt:

    tst = []
    x = data[:,0]
    y = data[:,1]  
    intt = data[:,2]
    total = []
    for i in range(1,len(data)):
       if intt[i]!=0:         
          tst.append((x[i]**2.0+ y[i]-y[i-1])**2.0)
       elif i == len(data) - 1:
          total.append(sum(tst))
       else:
          total.append(sum(tst))
          tst = []
          continue
    total = np.array(total)