简体   繁体   中英

Sum columns of parts of 2D array Python

how do you take the sum of the two last columns if the two first columns matches?

input:

 M =   [[1,1,3,5],
        [1,1,4,6],
        [1,2,3,7],
        [1,2,6,6],
        [2,1,0,8],
        [2,1,3,5],
        [2,2,9,6],
        [2,2,3,4]]

output:

M = [[1,1,7,11],
    [1,2,9,13],
    [2,1,3,13],
    [2,2,12,10]]

and can you do it whit a for loop?

Assuming that 2 similar lists follow each other always you could iterate over M[:-1] and then check the current list values against the next list values

M =[[1,1,3,5],[1,1,4,6],[1,2,3,7],[1,2,6,6],[2,1,0,8],[2,1,3,5],[2,2,9,6],[2,2,3,4]]
t=[]
for i,m in enumerate(M[:-1]):
  if m[0] == M[i+1][0] and m[1]==M[i+1][1]:
    t.append([m[0],m[1],m[2]+M[i+1][2],m[3]+M[i+1][3]])
print(t)
#[[1, 1, 7, 11], [1, 2, 9, 13], [2, 1, 3, 13], [2, 2, 12, 10]]

If the order might be scrambled I'd use 2 for loops. The second one will check m against every other list after it (it doesn't need to check those earlier since they checked against it).

for i,m in enumerate(M[:-1]):
  for x,n in enumerate(M[i+1:]):
    if m[0] == n[0] and m[1]==n[1]:
      t.append([m[0],m[1],m[2]+n[2],m[3]+n[3]])

We can find the unique tuples in the first two columns and then iterate over those to find the sum of each column whose rows equal the tuple.

Not sure what the fastest solution is, but this is one option:

M =[[1,1,3,5],[1,1,4,6],[1,2,3,7],[1,2,6,6],[2,1,0,8],[2,1,3,5],[2,2,9,6],[2,2,3,4]]

ans = []
for vals in list(set((x[0], x[1]) for x in M)):
    ans.append([vals[0], vals[1], sum(res[2] for res in M if (res[0], res[1]) == vals), sum(res[3] for res in M if (res[0], res[1]) == vals)])

A solution with list comprehension and itertools ' groupby :

from itertools import groupby

M = [
    [1,1,3,5],
    [1,1,4,6],
    [1,2,3,7],
    [1,2,6,6],
    [2,1,0,8],
    [2,1,3,5],
    [2,2,9,6],
    [2,2,3,4],
]

print([
    [
        key[0],
        key[1],
        sum(x[2] for x in group),
        sum(x[3] for x in group),
    ]
    for key, group in [
        (key, list(group))
        for key, group in groupby(sorted(M), lambda x: (x[0], x[1]))
    ]
])

Result:

[[1, 1, 7, 11], [1, 2, 9, 13], [2, 1, 3, 13], [2, 2, 12, 10]]

With reduce it can be simplified to:

from itertools import groupby
from functools import reduce

M = [
    [1,1,3,5],
    [1,1,4,6],
    [1,2,3,7],
    [1,2,6,6],
    [2,1,0,8],
    [2,1,3,5],
    [2,2,9,6],
    [2,2,3,4],
]

print([
    reduce(
        lambda x, y: [y[0], y[1], y[2] + x[2], y[3] + x[3]],
        group,
        (0, 0, 0, 0),
    )
    for _, group in groupby(sorted(M), lambda x: (x[0], x[1]))
])

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM